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}")