handler.py

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 inputs is already a string, try to parse it
           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."})
           
           # Generate the summary
           summary = self.summarizer(input_text)[0]["summary_text"]
           
           return json.dumps({"summary": summary})
       
       except Exception as e:
           return json.dumps({"error": str(e)})