| from typing import Dict |
| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| import os |
|
|
| class EndpointHandler: |
| """Custom handler for NuExtract-2-8B (InternLM2 based).""" |
|
|
| def __init__(self, path: str = "") -> None: |
| |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| path, |
| trust_remote_code=True |
| ) |
| self.model = AutoModelForCausalLM.from_pretrained( |
| path, |
| trust_remote_code=True, |
| torch_dtype=torch.float16, |
| device_map="auto" |
| ).eval() |
|
|
| def __call__(self, data: Dict[str, str]) -> Dict[str, str]: |
| prompt = data.get("inputs", "") |
| if not prompt: |
| return {"error": "No input provided."} |
|
|
| inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device) |
| output_ids = self.model.generate(**inputs, max_new_tokens=128) |
| answer = self.tokenizer.decode(output_ids[0], skip_special_tokens=True) |
| return {"generated_text": answer} |
