| import json |
| from transformers import pipeline |
|
|
| class EndpointHandler: |
| def __init__(self, model_dir): |
| """ |
| Initialize the handler with optimized settings for low-resource environments. |
| Args: |
| model_dir (str): The directory where the fine-tuned model is stored. |
| """ |
| print("Loading summarization pipeline...") |
| self.summarizer = pipeline( |
| "summarization", |
| model=model_dir, |
| device_map='cpu', |
| max_length=200, |
| min_length=60 |
| ) |
|
|
| def __call__(self, inputs): |
| try: |
| |
| if isinstance(inputs, str): |
| try: |
| input_data = json.loads(inputs) |
| except json.JSONDecodeError: |
| input_data = {"inputs": inputs} |
| elif isinstance(inputs, dict): |
| input_data = inputs |
| else: |
| raise ValueError("Input must be a string or dictionary") |
| |
| input_text = input_data.get("inputs", "") |
| if not input_text: |
| return json.dumps({"error": "No input text provided."}) |
| |
| |
| summary = self.summarizer(input_text)[0]["summary_text"] |
| |
| return json.dumps({"summary": summary}) |
| |
| except Exception as e: |
| return json.dumps({"error": str(e)}) |
