We are trying to host a hugging face model Gliner for NER on vertex using custom prediction routine . Below is the code . Issue is that uvicorn server starts parent process but it shut down. We dont know the reason . We even tried putting logs but they are note getting printed . Please guide .
from google.cloud.aiplatform.prediction.predictor import Predictor
from google.cloud.aiplatform.utils import prediction_utils
import torch
from gliner import GLiNER
from typing import Dict
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
labels = [
"person",
"location",
"award",
"date",
"company",
"organization",
"product",
"technology",
"law",
"credit card number",
"job title",
]
class GlinerPredictor(Predictor😞
def __init__(self😞
# self._class_names = ["setosa", "versicolor", "virginica"]
pass
def load(self, artifacts_uri: str😞
"""Loads the model artifacts."""
prediction_utils.download_model_artifacts(artifacts_uri)
print("after downloading model")
logger.info("after downloading model")
self._model = GLiNER.from_pretrained(
"urchade/gliner_large-v2.1"
) # torch.load("model.pt")
print("after loading model")
logger.info("after loading model")
def preprocess(self, prediction_input: Dict) -> Dict:
return prediction_input
def predict(self, prediction_input: Dict) -> Dict:
"""Performs prediction."""
logger.info(f"prediction_input {prediction_input}")
instances = prediction_input["instances"]
outputs = []
for instance in instances:
entities = self._model.predict_entities(instance["prompt"], labels)
output = ""
for entity in entities:
# print(entity["text"], "=>", entity["label"])
output += entity["text"] + "=>" + entity["label"]
output += "\n"
outputs.append(output)
prediction_results = {"predictions": [o for o in outputs]}
return prediction_results
def postprocess(self, prediction_results: Dict) -> Dict:
return prediction_results