Update app.py

app.py

@@ -1,93 +1,85 @@
-# app.py
-
-import json
-import re
-import torch
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import torch
+import json
+import re
+
 from huggingface_hub import snapshot_download
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import PeftModel
 
-# -----------------------------
-# Load Model at Startup
-# -----------------------------
+app = FastAPI(title="CV–Job Description Matching API")
+
+# ---------- Request body ----------
+class MatchRequest(BaseModel):
+    cv: str
+    job_description: str
 
+# ---------- Load model once ----------
 BASE_MODEL = "akjindal53244/Llama-3.1-Storm-8B"
 ADAPTER_MODEL = "LlamaFactoryAI/cv-job-description-matching"
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-)
+model = None
+tokenizer = None
 
-print("Downloading adapter...")
-adapter_path = snapshot_download(ADAPTER_MODEL)
-
-# Patch adapter_config.json exactly like in Kaggle
-config_path = adapter_path + "/adapter_config.json"
-with open(config_path, "r") as f:
-    cfg = json.load(f)
-cfg["task_type"] = "CAUSAL_LM"
-with open(config_path, "w") as f:
-    json.dump(cfg, f, indent=2)
-print("Patched adapter_config.json")
-
-print("Loading tokenizer + base model...")
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    quantization_config=bnb_config,
-    device_map="auto",
-)
-base_model.config.pad_token_id = tokenizer.pad_token_id
+def load_model():
+    global model, tokenizer
 
-print("Applying LoRA adapter...")
-model = PeftModel.from_pretrained(
-    base_model,
-    adapter_path,
-    device_map="auto"
-)
-model.eval()
-torch.set_grad_enabled(False)
+    if model is not None:
+        return
 
-print("Model ready.")
+    print("Downloading adapter...")
+    adapter_path = snapshot_download(ADAPTER_MODEL)
 
-# -----------------------------
-# FastAPI Setup
-# -----------------------------
+    # Patch adapter_config.json
+    cfg_path = adapter_path + "/adapter_config.json"
+    with open(cfg_path, "r") as f:
+        cfg = json.load(f)
+    cfg["task_type"] = "CAUSAL_LM"
+    with open(cfg_path, "w") as f:
+        json.dump(cfg, f, indent=2)
 
-app = FastAPI()
+    print("Loading tokenizer & base model...")
+    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
 
-class MatchRequest(BaseModel):
-    cv: str
-    job_description: str
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
 
-@app.get("/")
-def root():
-    return {"status": "ok", "message": "CV Matching API running"}
+    base = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        quantization_config=bnb,
+        device_map="auto",
+    )
+    base.config.pad_token_id = tokenizer.pad_token_id
+
+    print("Loading LoRA adapter...")
+    model = PeftModel.from_pretrained(base, adapter_path, device_map="auto")
+    model.eval()
+    torch.set_grad_enabled(False)
+
+    print("Model is ready.")
+
+@app.on_event("startup")
+def startup_event():
+    load_model()
+
+# ---------- System prompt ----------
+SYSTEM_PROMPT = (
+    "You analyze how well a CV matches a job description. "
+    "Your ONLY output must be JSON with keys: "
+    "matching_analysis, description, score, recommendation."
+)
+
+# ---------- Run inference ----------
+def run_inference(cv, jd):
+    global model, tokenizer
 
-@app.post("/predict")
-def predict(req: MatchRequest):
     messages = [
-        {
-            "role": "system",
-            "content": (
-                "You analyze how well a CV matches a job description. "
-                "Your ONLY output must be JSON with the keys: "
-                "matching_analysis, description, score, recommendation."
-            ),
-        },
-        {
-            "role": "user",
-            "content": f"<CV> {req.cv} </CV>\n<job_description> {req.job_description} </job_description>",
-        },
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": f"<CV> {cv} </CV><job_description> {jd} </job_description>"}
     ]
 
-    # Build chat prompt
     prompt = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
@@ -98,20 +90,27 @@ def predict(req: MatchRequest):
     encoded = {k: v.to(model.device) for k, v in encoded.items()}
 
     with torch.inference_mode():
-        output = model.generate(
+        out = model.generate(
             **encoded,
             max_new_tokens=256,
             pad_token_id=tokenizer.pad_token_id,
         )
 
     input_len = encoded["input_ids"].shape[1]
-    generated = tokenizer.decode(output[0][input_len:], skip_special_tokens=True)
-
-    # Try to parse JSON
-    try:
-        start = generated.index("{")
-        end = generated.rindex("}") + 1
-        json_text = generated[start:end]
-        return json.loads(json_text)
-    except Exception:
-        return {"raw_output": generated}
+    generated = tokenizer.decode(out[0][input_len:], skip_special_tokens=True)
+
+    # Extract JSON
+    match = re.search(r"\{.*\}", generated, re.DOTALL)
+    if match:
+        return json.loads(match.group(0))
+
+    return {"raw_output": generated}
+
+# ---------- API route ----------
+@app.post("/match")
+def match(request: MatchRequest):
+    return run_inference(request.cv, request.job_description)
+
+@app.get("/")
+def root():
+    return {"message": "API running. POST /match to use it."}