| |
| import os, json, time, tempfile, base64 |
| import torch |
| from docling.document_converter import DocumentConverter |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from supabase import create_client |
| from huggingface_hub import EndpointHandler |
|
|
| class Handler(EndpointHandler): |
| def __init__(self): |
| from dotenv import load_dotenv |
| load_dotenv() |
| |
| self.supabase = create_client( |
| os.getenv("DATABASE_URL"), |
| os.getenv("SUPABASE_SERVICE_ROLE_KEY") |
| ) |
| |
| device = "mps" if torch.backends.mps.is_available() else "cpu" |
| dtype = torch.float16 if device=="mps" else torch.float32 |
| self.model = AutoModelForCausalLM.from_pretrained( |
| "numind/NuExtract-1.5-tiny", |
| torch_dtype=dtype, trust_remote_code=True |
| ).to(device).eval() |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| "numind/NuExtract-1.5-tiny", trust_remote_code=True |
| ) |
|
|
| def __call__(self, payload): |
| """ |
| Expects JSON: { |
| "inputs": "<raw text or base64-pdf>", |
| "is_pdf": false|true |
| } |
| """ |
| text = payload["inputs"] |
| if payload.get("is_pdf"): |
| b = base64.b64decode(text) |
| with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: |
| tmp.write(b) |
| path = tmp.name |
| conv = DocumentConverter().convert(path) |
| text = conv.document.export_to_text() |
| os.remove(path) |
|
|
| |
| prompt = ( |
| "<|input|>\n### Instruction:\nRemplis la template JSON…\n" |
| "### Text:\n" + text + "\n<|output|>" |
| ) |
| enc = self.tokenizer(prompt, return_tensors="pt", |
| truncation=True).to(self.model.device) |
| out_ids = self.model.generate(**enc, max_new_tokens=1024) |
| out = self.tokenizer.decode(out_ids[0], skip_special_tokens=True) |
| result = out.split("<|output|>")[1] |
| return {"result": json.loads(result)} |
