Marcin-XStudio/CV_Exractor

1import json2import torch3import time4from docling.document_converter import DocumentConverter5from transformers import AutoModelForCausalLM, AutoTokenizer6from pathlib import Path7from fastapi import FastAPI, File, UploadFile, HTTPException8import os9from dotenv import load_dotenv10import tempfile11from supabase import create_client12from huggingface_hub import snapshot_download13from transformers import BitsAndBytesConfig, AutoModelForCausalLM14 15load_dotenv()16app = FastAPI()17 18model_name = "numind/NuExtract-1.5"19 20# MODEL_CACHE = "/home/user/app/model_cache"21 22device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"23dtype = torch.float16 if device in ("mps", "cuda") else torch.float3224 25print("CUDA available:", torch.cuda.is_available())       # True26print("Device name:", torch.cuda.get_device_name(0)) 27print ("Model Running ", model_name)28 29# If lower memory usage needed: 30 31# bnb_config = BitsAndBytesConfig(32#     load_in_4bit=True,33#     bnb_4bit_use_double_quant=True,34#     bnb_4bit_quant_type="nf4",35#     bnb_4bit_compute_dtype=torch.float1636# )37 38@app.on_event("startup")39def startup_supabase():40    print("DEVICE:", device)41    global supabase42    supabase = create_client(43        os.getenv("DATABASE_URL"),44        os.getenv("SUPABASE_SERVICE_ROLE_KEY")45    )46 47@app.on_event("startup")48def load_model():49    print("Loading model and tokenizer...", flush=True)50    global model, tokenizer51    model = AutoModelForCausalLM.from_pretrained(52        # model_name,53        # cache_dir=MODEL_CACHE,54        MODEL_CACHE,55        local_files_only=True,56        torch_dtype=dtype,57        trust_remote_code=True,58        # quantization_config=bnb_config,59        device_map="auto"60    ).to(device).eval()61    tokenizer = AutoTokenizer.from_pretrained(62        # model_name,63        # cache_dir=MODEL_CACHE,64        MODEL_CACHE,65        local_files_only=True,66        trust_remote_code=True,67        device_map="auto"68    )69    # Check this optimization!70    # if torch.__version__ >= "2.0":71    #     model = torch.compile(model)72    print("✅ Model and tokenizer loaded from", MODEL_CACHE)73 74 75def predict_NuExtract(texts, template, batch_size=1, max_length=5096, max_new_tokens=1024):76    print("Starting NuExtract prediction...", flush=True)77    start_time = time.perf_counter()78    template_str = json.dumps(json.loads(template), indent=4)79    prompts = [80        "<|input|>\n"81        "### Instruction:\n"82        "Remplis la template JSON avec les informations extraits du texte.\n"83        "Extraire le nom du candidat tel qu’il apparaît sur la première ligne du document (souvent en majuscules ou en plus gros), et le mettre dans nom. Si le nom n’est pas trouvé, renvoyer une chaîne vide.\n"84        "Si le text contient des mentions de diplomes ou de titres de formations, certificats ou Attestation on les considère comme education, pas experience\n"85        "Pour chaque bloc formation ou expérience où une date est mentionnée (MM/YYYY ou mois/YYYY etc), remplir systématiquement annee_debut et annee_fin.\n"86        "Ne jamais laisser ces champs vides si la date est dans le texte.\n"87        "Si une seule date specifié pour une experience ou une formation met la même date pour start_date et end_date. Exemples : \n"88        "Exemple 1 : \n"89        "Texte : \"...01/2024 – En cours...\"\n"90        "Output : \"start_date\": \"01/2024\", \"end_date\": \"01/2024\"\n"91        "Exemple 2 : \n"92        "Texte : \"...mai 2025...\"\n"93        "Output : \"start_date\": \"05/2025\", \"end_date\": \"05/2025\"\n"94        "Exemple 3 : \n"95        "Texte : \"...mai 2025 – juin 2026...\"\n"96        "Output : \"start_date\": \"05/2025\", \"end_date\": \"06/2026\"\n"97        "Extraction des dates est *très* importante. Ne laisse jamais les dates vides\n"98        "Exemples types de formations :  CAP Boucherie, Licence Pro Métiers de l’Énergétique, Baccalauréat Général\n"99        "Exemples catégories de formations :  Transport, énergie, langues, esthétique\n"100        "Exemples mobilités :  permis B, permis C, permis D. si y'a juste la mention de permis on considère que c'est le permis B. N’inclure que les permis explicitement mentionnés dans le texte. S’il n’y a aucune mention de permis, renvoyer une liste vide.\n"101        "### Exemple 1 (avec permis) \n"102        "Texte : \"...j’ai obtenu mon permis B en 2015...\"\n"103        "Output : \"types_des_permis_de_conduire\": [\"permis B\"]\n"104        "### Exemple 2 (sans permis) \n"105        "Texte : \"...j’ai étudié à l’Université de Paris...\"\n"106        "Output : \"types_des_permis_de_conduire\": []\n"107        "Output *only* the completed JSON.\n"108        "### Template:\n"109        f"{template_str}\n"110        "### Text:\n"111        f"{text}\n\n"112        "<|output|>"113        for text in texts114    ]115    print("Prompts prepared.", flush=True)116    outputs = []117    with torch.no_grad():118        for i in range(0, len(prompts), batch_size):119            batch = prompts[i : i+batch_size]120            enc = tokenizer(121                batch,122                return_tensors="pt",123                truncation=True,124                padding=True,125                max_length=max_length126            ).to(device)127            print(f"Generating outputs with model for batch {i//batch_size+1}...", flush=True)128            ids = model.generate(**enc, max_new_tokens=max_new_tokens, num_beams=1, use_cache=False)129            outputs += tokenizer.batch_decode(ids, skip_special_tokens=True)130    print("Outputs generated.", flush=True)131    elapsed = time.perf_counter() - start_time132    print(f"NuExtract prediction completed in {elapsed:.2f} seconds.", flush=True)133    return [out.split("<|output|>")[1] for out in outputs]134 135template = """{136    "nom": "", "email": "", "telephone": "",137    "education": [{"type_de_formation": "", "categorie_de_formation": "", "annee_debut": "", "annee_fin": ""}],138    "experience": [{"position": "", "entreprise": "", "annee_debut": "", "annee_fin": ""}],139    "types_des_permis_de_conduire": [""]140}"""141 142data_model = {143        "experience": [{"start_date": "", "end_date": "", "job_category_id": ""}],144        "education": [{"training_type_id": "", "training_category_id": "", "start_date": "", "end_date": ""}],145        "email": "",146        "phone": "",147        "mobility": [{"id": "", "title": ""}]148}149 150@app.get("/health")151async def health():152    return {"status": "ok"}153 154@app.post("/extract")155async def extract(file: UploadFile = File(...)):156    suffix = Path(file.filename).suffix or ".pdf"157    try:158        # Create one global client; reused across calls159        # supabase = create_client(os.environ.get("DATABASE_URL"), os.environ.get("SUPABASE_SERVICE_ROLE_KEY"))160 161        # Use Supabase client to query tables, map names to IDs162        job_categories = supabase.table("Job_category").select("id, title").execute()163 164        print("Job Categories: ", job_categories)165 166        training_types = supabase.table("Training_type").select("id, title").execute()167 168        print("Training Types: ", training_types)169 170        training_categories = supabase.table("Training_category").select("id, title").execute()171 172        print("Training Categories: ", training_categories)173 174        mobility = supabase.table("Mobility").select("id, title").execute()175 176        print("Mobility: ", mobility)177 178        # … your LLM logic …179            180        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:181            data = await file.read()182            tmp.write(data)183            tmp_path = tmp.name184        print(f"Upload saved to {tmp_path}", flush=True)185    except Exception as e:186        print(f"Cannot save upload: {e}", flush=True)187        raise HTTPException(400, f"Cannot save upload: {e}")188 189    try:190        converter = DocumentConverter()191        result    = converter.convert(tmp_path)192        raw_text  = result.document.export_to_text()193        print("Docling conversion complete.", flush=True)194    except Exception as e:195        print(f"Docling error: {e}", flush=True)196        raise HTTPException(500, f"Docling error: {e}")197    finally:198        try: os.remove(tmp_path)199        except OSError: pass200 201    try:202        extracted_json_str = predict_NuExtract([raw_text], template)[0]203        print("Extraction with NuExtract complete.", flush=True)204        print("⏺ RAW MODEL OUTPUT:\n", extracted_json_str)205        print("⏺ RAW MODEL OUTPUT (repr):\n", repr(extracted_json_str))206 207        print("Clearing Cache", flush=True)208        if device == "mps":209            torch.mps.empty_cache()210        elif device == "cuda":211            torch.cuda.empty_cache()212        elif device == "cpu":213            torch.cpu.empty_cache()214        return {"result": json.loads(extracted_json_str)}215    except Exception as e:216        print(f"Extraction error: {e}", flush=True)217        raise HTTPException(500, f"Extraction error: {e}")