File size: 1,400 Bytes
3dd9480 576cae6 3dd9480 576cae6 3dd9480 576cae6 3dd9480 576cae6 |
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 |
from typing import Dict, List, Any
from transformers import pipeline, AutoTokenizer, BartForConditionalGeneration
class EndpointHandler():
def __init__(self, path=""):
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
self.model = BartForConditionalGeneration.from_pretrained(path).to(self.device)
self.tokenizer = AutoTokenizer.from_pretrained(path)
except Exception as e:
print(f"Error loading model or tokenizer from path {path}: {e}")
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
inputs (:obj: `str`)
date (:obj: `str`)
Return:
A :obj:`list` | `dict`: will be serialized and returned
"""
# get inputs
inputs = data.get("inputs", "")
if not inputs:
return [{"error": "No inputs provided"}]
tokenized_input = self.tokenizer(inputs, return_tensors="pt", truncation=True, max_length=1024, padding="max_length")
tokenized_input = tokenized_input.to(self.device) # Move input tensors to the same device as model
summary_ids = self.model.generate(**tokenized_input, max_length=256, do_sample=True, top_p=0.8)
summary_text = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return [{"summary": summary_text}] |