Hhhh / summarization_api.py
Kfjjdjdjdhdhd's picture
Upload 26 files
7b74407 verified
from flask import jsonify, send_file, request
from main import *
import torch, io, base64
def summarize_text(text, output_path="output_summary.txt"):
if summarization_model is None or summarization_tokenizer is None: return {"error": "Summarization model or tokenizer not initialized."}
input_ids = summarization_tokenizer.encode(text, return_tensors="pt").to(device)
with torch.no_grad(): summary_ids = summarization_model.generate(input_ids, num_beams=4, max_length=100, early_stopping=True); summary_text = summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return {"summary_text": summary_text}
def summarization_api(text):
output = summarize_text(text)
if "error" in output: return {"error": output["error"]}
return output