File size: 1,420 Bytes
72d83da | 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 | # from gliner import GLiNER
# class EndpointHandler:
# def __init__(self, path=""):
# # Use the provided path for loading the model
# self.model = GLiNER.from_pretrained(path)
# def __call__(self, data):
# try:
# text = data.get("text", "")
# labels = data.get("labels", [])
# if not text or not labels:
# return {"error": "Please provide 'text' and 'labels'"}
# entities = self.model.predict_entities(text, labels)
# return {"entities": entities}
# except Exception as e:
# return {"error": str(e)}
from gliner import GLiNER
import torch
class EndpointHandler:
def __init__(self, path=""):
# Load without device_map, then move to GPU
self.model = GLiNER.from_pretrained(path) # Remove device_map="cuda"
self.model = self.model.to("cuda")
self.model.eval() # Lock for inference
def __call__(self, data):
# If data is wrapped in 'inputs' (as Hugging Face does), unwrap it
if isinstance(data, dict) and "inputs" in data:
data = data["inputs"]
text = data.get("text", "")
labels = data.get("labels", [])
if not text or not labels:
return {"error": "Please provide 'text' and 'labels'"}
entities = self.model.predict_entities(text, labels)
return {"entities": entities}
|