| 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 |
| """ |
| |
| |
|
|
| 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) |
|
|
| 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}] |
