File size: 2,118 Bytes
f439d97 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | # handler.py
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()
# supabase
self.supabase = create_client(
os.getenv("DATABASE_URL"),
os.getenv("SUPABASE_SERVICE_ROLE_KEY")
)
# model
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)
# build your prompt + generate (same as predict_NuExtract)
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)} |