Update README.md

README.md

@@ -72,3 +72,67 @@ base_model = AutoModelForCausalLM.from_pretrained(
 model = PeftModel.from_pretrained(base_model, lora_model_id, device_map="auto")
 model = model.merge_and_unload()
 model.eval()
+
+def extract_and_clean_json(text):
+    """Extract JSON from LLM output, even if extra text is present."""
+    match = re.search(r"\{[\s\S]*\}", text)
+    if not match:
+        return None
+
+    json_str = match.group(0)
+    json_str = json_str.replace("None", "null")
+    json_str = json_str.replace("True", "true").replace("False", "false")
+    json_str = re.sub(r",(\s*[}\]])", r"\1", json_str)
+
+    try:
+        return json5.loads(json_str)
+    except Exception as e:
+        print(f"JSON parse error: {e}")
+        return None
+
+
+def infer_from_text(jd_text: str):
+    """Runs inference on a job description."""
+    start_time = time.time()
+
+    system_prompt = """Extract structured information from the following job description and return it as JSON.
+""".strip()
+
+    user_prompt = f"""
+Job Description:
+{jd_text}
+""".strip()
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt}
+    ]
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    raw_inputs = tokenizer(prompt, return_tensors="pt")
+    device = model.device
+    inputs = {k: v.to(device) for k, v in raw_inputs.items()}
+
+    with torch.no_grad():
+        out = model.generate(
+            **inputs,
+            max_new_tokens=1000,
+            do_sample=False,
+            temperature=0,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    gen_tokens = out[0][inputs["input_ids"].shape[1]:]
+    response_text = tokenizer.decode(gen_tokens, skip_special_tokens=True)
+    duration = round(time.time() - start_time, 2)
+
+    parsed = extract_and_clean_json(response_text)
+    if parsed is not None:
+        return json.dumps(parsed, indent=2), duration
+
+    return response_text, duration