File size: 1,439 Bytes
11a0afe f210283 584224e 869c622 584224e 11a0afe 584224e |
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 37 38 39 40 41 42 43 |
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)}) |