app.py

import json
import torch
import time
from docling.document_converter import DocumentConverter
from transformers import AutoModelForCausalLM, AutoTokenizer
from pathlib import Path
from fastapi import FastAPI, File, UploadFile, HTTPException
import os
from dotenv import load_dotenv
import tempfile
from supabase import create_client
from huggingface_hub import snapshot_download
from transformers import BitsAndBytesConfig, AutoModelForCausalLM

load_dotenv()
app = FastAPI()

model_name = "numind/NuExtract-1.5"

# MODEL_CACHE = "/home/user/app/model_cache"

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
dtype = torch.float16 if device in ("mps", "cuda") else torch.float32

print("CUDA available:", torch.cuda.is_available())       # True
print("Device name:", torch.cuda.get_device_name(0)) 
print ("Model Running ", model_name)

# If lower memory usage needed: 

# bnb_config = BitsAndBytesConfig(
#     load_in_4bit=True,
#     bnb_4bit_use_double_quant=True,
#     bnb_4bit_quant_type="nf4",
#     bnb_4bit_compute_dtype=torch.float16
# )

@app.on_event("startup")
def startup_supabase():
    print("DEVICE:", device)
    global supabase
    supabase = create_client(
        os.getenv("DATABASE_URL"),
        os.getenv("SUPABASE_SERVICE_ROLE_KEY")
    )

@app.on_event("startup")
def load_model():
    print("Loading model and tokenizer...", flush=True)
    global model, tokenizer
    model = AutoModelForCausalLM.from_pretrained(
        # model_name,
        # cache_dir=MODEL_CACHE,
        MODEL_CACHE,
        local_files_only=True,
        torch_dtype=dtype,
        trust_remote_code=True,
        # quantization_config=bnb_config,
        device_map="auto"
    ).to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(
        # model_name,
        # cache_dir=MODEL_CACHE,
        MODEL_CACHE,
        local_files_only=True,
        trust_remote_code=True,
        device_map="auto"
    )
    # Check this optimization!
    # if torch.__version__ >= "2.0":
    #     model = torch.compile(model)
    print("✅ Model and tokenizer loaded from", MODEL_CACHE)


def predict_NuExtract(texts, template, batch_size=1, max_length=5096, max_new_tokens=1024):
    print("Starting NuExtract prediction...", flush=True)
    start_time = time.perf_counter()
    template_str = json.dumps(json.loads(template), indent=4)
    prompts = [
        "<|input|>\n"
        "### Instruction:\n"
        "Remplis la template JSON avec les informations extraits du texte.\n"
        "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"
        "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"
        "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"
        "Ne jamais laisser ces champs vides si la date est dans le texte.\n"
        "Si une seule date specifié pour une experience ou une formation met la même date pour start_date et end_date. Exemples : \n"
        "Exemple 1 : \n"
        "Texte : \"...01/2024 – En cours...\"\n"
        "Output : \"start_date\": \"01/2024\", \"end_date\": \"01/2024\"\n"
        "Exemple 2 : \n"
        "Texte : \"...mai 2025...\"\n"
        "Output : \"start_date\": \"05/2025\", \"end_date\": \"05/2025\"\n"
        "Exemple 3 : \n"
        "Texte : \"...mai 2025 – juin 2026...\"\n"
        "Output : \"start_date\": \"05/2025\", \"end_date\": \"06/2026\"\n"
        "Extraction des dates est *très* importante. Ne laisse jamais les dates vides\n"
        "Exemples types de formations :  CAP Boucherie, Licence Pro Métiers de l’Énergétique, Baccalauréat Général\n"
        "Exemples catégories de formations :  Transport, énergie, langues, esthétique\n"
        "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"
        "### Exemple 1 (avec permis) \n"
        "Texte : \"...j’ai obtenu mon permis B en 2015...\"\n"
        "Output : \"types_des_permis_de_conduire\": [\"permis B\"]\n"
        "### Exemple 2 (sans permis) \n"
        "Texte : \"...j’ai étudié à l’Université de Paris...\"\n"
        "Output : \"types_des_permis_de_conduire\": []\n"
        "Output *only* the completed JSON.\n"
        "### Template:\n"
        f"{template_str}\n"
        "### Text:\n"
        f"{text}\n\n"
        "<|output|>"
        for text in texts
    ]
    print("Prompts prepared.", flush=True)
    outputs = []
    with torch.no_grad():
        for i in range(0, len(prompts), batch_size):
            batch = prompts[i : i+batch_size]
            enc = tokenizer(
                batch,
                return_tensors="pt",
                truncation=True,
                padding=True,
                max_length=max_length
            ).to(device)
            print(f"Generating outputs with model for batch {i//batch_size+1}...", flush=True)
            ids = model.generate(**enc, max_new_tokens=max_new_tokens, num_beams=1, use_cache=False)
            outputs += tokenizer.batch_decode(ids, skip_special_tokens=True)
    print("Outputs generated.", flush=True)
    elapsed = time.perf_counter() - start_time
    print(f"NuExtract prediction completed in {elapsed:.2f} seconds.", flush=True)
    return [out.split("<|output|>")[1] for out in outputs]

template = """{
    "nom": "", "email": "", "telephone": "",
    "education": [{"type_de_formation": "", "categorie_de_formation": "", "annee_debut": "", "annee_fin": ""}],
    "experience": [{"position": "", "entreprise": "", "annee_debut": "", "annee_fin": ""}],
    "types_des_permis_de_conduire": [""]
}"""

data_model = {
        "experience": [{"start_date": "", "end_date": "", "job_category_id": ""}],
        "education": [{"training_type_id": "", "training_category_id": "", "start_date": "", "end_date": ""}],
        "email": "",
        "phone": "",
        "mobility": [{"id": "", "title": ""}]
}

@app.get("/health")
async def health():
    return {"status": "ok"}

@app.post("/extract")
async def extract(file: UploadFile = File(...)):
    suffix = Path(file.filename).suffix or ".pdf"
    try:
        # Create one global client; reused across calls
        # supabase = create_client(os.environ.get("DATABASE_URL"), os.environ.get("SUPABASE_SERVICE_ROLE_KEY"))

        # Use Supabase client to query tables, map names to IDs
        job_categories = supabase.table("Job_category").select("id, title").execute()

        print("Job Categories: ", job_categories)

        training_types = supabase.table("Training_type").select("id, title").execute()

        print("Training Types: ", training_types)

        training_categories = supabase.table("Training_category").select("id, title").execute()

        print("Training Categories: ", training_categories)

        mobility = supabase.table("Mobility").select("id, title").execute()

        print("Mobility: ", mobility)

        # … your LLM logic …
            
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            data = await file.read()
            tmp.write(data)
            tmp_path = tmp.name
        print(f"Upload saved to {tmp_path}", flush=True)
    except Exception as e:
        print(f"Cannot save upload: {e}", flush=True)
        raise HTTPException(400, f"Cannot save upload: {e}")

    try:
        converter = DocumentConverter()
        result    = converter.convert(tmp_path)
        raw_text  = result.document.export_to_text()
        print("Docling conversion complete.", flush=True)
    except Exception as e:
        print(f"Docling error: {e}", flush=True)
        raise HTTPException(500, f"Docling error: {e}")
    finally:
        try: os.remove(tmp_path)
        except OSError: pass

    try:
        extracted_json_str = predict_NuExtract([raw_text], template)[0]
        print("Extraction with NuExtract complete.", flush=True)
        print("⏺ RAW MODEL OUTPUT:\n", extracted_json_str)
        print("⏺ RAW MODEL OUTPUT (repr):\n", repr(extracted_json_str))

        print("Clearing Cache", flush=True)
        if device == "mps":
            torch.mps.empty_cache()
        elif device == "cuda":
            torch.cuda.empty_cache()
        elif device == "cpu":
            torch.cpu.empty_cache()
        return {"result": json.loads(extracted_json_str)}
    except Exception as e:
        print(f"Extraction error: {e}", flush=True)
        raise HTTPException(500, f"Extraction error: {e}")