server.py

import json
import re
import os
from typing import Any, Dict, Optional

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModelForCausalLM


REPO_ID = "Alibaba-EI/SmartResume"
SUBFOLDER = "Qwen3-0.6B"

LOCAL_DIR = os.getenv("MODEL_CACHE_DIR", "./models/smartresume")
MODEL_PATH = f"{LOCAL_DIR}/{SUBFOLDER}"

# A40: tu peux monter le contexte sans stress, mais reste pragmatique
MAX_MODEL_LEN = int(os.getenv("MAX_MODEL_LEN", "8192"))
MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "1200"))

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32


# --------- Helpers ---------

def download_model_if_needed() -> None:
    if os.path.isdir(MODEL_PATH) and os.path.exists(os.path.join(MODEL_PATH, "config.json")):
        return
    snapshot_download(
        repo_id=REPO_ID,
        local_dir=LOCAL_DIR,
        local_dir_use_symlinks=False,
    )

def build_schema() -> Dict[str, Any]:
    return {
        "basics": {
            "full_name": "",
            "email": "",
            "phone": "",
            "location": "",
            "links": []
        },
        "headline": "",
        "summary": "",
        "skills": [],
        "experience": [
            {
                "company": "",
                "title": "",
                "start_date": "",
                "end_date": "",
                "location": "",
                "highlights": []
            }
        ],
        "education": [
            {
                "institution": "",
                "degree": "",
                "field": "",
                "start_date": "",
                "end_date": ""
            }
        ],
        "certifications": [],
        "languages": [],
        "projects": [],
        "publications": []
    }

def extract_json(text: str) -> Dict[str, Any]:
    text = text.strip()
    # direct
    try:
        return json.loads(text)
    except Exception:
        pass

    # find largest {...}
    candidates = re.findall(r"\{.*\}", text, flags=re.DOTALL)
    if not candidates:
        raise ValueError("No JSON object found in model output.")

    candidates.sort(key=len, reverse=True)
    for c in candidates:
        try:
            return json.loads(c)
        except Exception:
            continue

    raise ValueError("Found JSON-like text but invalid JSON.")

def build_prompt(tokenizer, cv_text: str, schema: Dict[str, Any]) -> str:
    system = (
        "You are a CV parsing engine.\n"
        "Return ONLY valid JSON (RFC8259). No markdown. No commentary.\n"
        "Do not invent facts. If missing, use empty string or empty list.\n"
        "Dates: YYYY-MM or YYYY-MM-DD when possible.\n"
    )
    user = (
        "Extract this CV into EXACTLY the following JSON shape.\n"
        "Output must be STRICT JSON only.\n\n"
        f"JSON SHAPE:\n{json.dumps(schema, ensure_ascii=False)}\n\n"
        f"CV TEXT:\n{cv_text}\n"
    )
    messages = [{"role": "system", "content": system},
                {"role": "user", "content": user}]

    if hasattr(tokenizer, "apply_chat_template"):
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    return system + "\n" + user + "\nJSON:\n"


# --------- FastAPI ---------

class ParseRequest(BaseModel):
    cv_text: str = Field(..., min_length=10)
    schema: Optional[Dict[str, Any]] = None  # allow custom schema per request
    max_new_tokens: Optional[int] = None
    temperature: float = 0.0

class ParseResponse(BaseModel):
    data: Dict[str, Any]
    model: str
    device: str


app = FastAPI(title="SmartResume CV Parser", version="1.0.0")

tokenizer = None
model = None

@app.on_event("startup")
def startup():
    global tokenizer, model
    download_model_if_needed()

    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)

    # bfloat16 sur A40: top
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=DTYPE,
        device_map="auto" if DEVICE == "cuda" else None,
        trust_remote_code=True
    )
    model.eval()

@app.get("/health")
def health():
    return {"status": "ok", "device": DEVICE, "model_path": MODEL_PATH}

@app.post("/parse", response_model=ParseResponse)
def parse(req: ParseRequest):
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    schema = req.schema or build_schema()
    prompt = build_prompt(tokenizer, req.cv_text, schema)

    # Tokenize
    inputs = tokenizer(prompt, return_tensors="pt")
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    max_new = req.max_new_tokens or MAX_NEW_TOKENS

    # Gen deterministic
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new,
            do_sample=False if req.temperature == 0 else True,
            temperature=req.temperature,
            top_p=1.0,
            eos_token_id=tokenizer.eos_token_id,
        )

    decoded = tokenizer.decode(out[0], skip_special_tokens=True)
    if decoded.startswith(prompt):
        decoded = decoded[len(prompt):].strip()

    # Parse JSON; if fail, one repair pass
    try:
        data = extract_json(decoded)
    except Exception:
        repair = (
            "Your previous output was invalid JSON.\n"
            "Fix it and return ONLY valid JSON. No markdown.\n\n"
            f"INVALID OUTPUT:\n{decoded}\n"
        )
        repair_inputs = tokenizer(repair, return_tensors="pt")
        repair_inputs = {k: v.to(model.device) for k, v in repair_inputs.items()}

        with torch.no_grad():
            out2 = model.generate(
                **repair_inputs,
                max_new_tokens=max_new,
                do_sample=False,
                temperature=0.0,
                top_p=1.0,
                eos_token_id=tokenizer.eos_token_id,
            )

        decoded2 = tokenizer.decode(out2[0], skip_special_tokens=True)
        data = extract_json(decoded2)

    return ParseResponse(data=data, model=MODEL_PATH, device=DEVICE)