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()


# // FOR RUNNING IN SPACES
model_name = "numind/NuExtract-1.5-tiny"
# Path inside your container
# MODEL_PATH = "/app/model_cache/models--numind--NuExtract-1.5-tiny/snapshots/df52efb3109d324cd52b30728f9e3fdedf19f742"
# If you used local_dir="model", snapshot_download will still create models--… subfolder.
# You can also symlink or copy it to /app/model directly in Dockerfile.


# MODEL_PATH = "/app/model_cache"
# model_cache_path = snapshot_download(
#     repo_id="numind/NuExtract-1.5-tiny",
#     local_dir="/app/model_cache",            # <-- direct destination
#     cache_dir="/app/model_cache/hf_cache"
# )

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

print(">>> MODEL CACHE PATH:", MODEL_CACHE, os.listdir(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)) 
# bnb_config = BitsAndBytesConfig(load_in_8bit=True)


# If lower memory usage needed: 

# bnb_config = BitsAndBytesConfig(
#     load_in_4bit=True,
#     bnb_4bit_use_double_quant=True,
#     bnb_4bit_quant_type="nf4"
# )
# model = AutoModelForCausalLM.from_pretrained(
#     MODEL_CACHE,
#     quantization_config=bnb_config,
#     device_map="auto",
#     local_files_only=True,
#     trust_remote_code=True
# )

@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, torch_dtype=dtype, trust_remote_code=True
    # )
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_CACHE,
        local_files_only=True,
        torch_dtype=dtype,
        trust_remote_code=True,
        # quantization_config=bnb_config,
        # no_split_module_classes=["Block"],
        device_map="auto"
    ).to(device).eval()
    # tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_CACHE,
        local_files_only=True,
        trust_remote_code=True,
        device_map="auto"
    )
    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"
        "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"
        "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, 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}")