Hhhh / summarization_api.py
Hjgugugjhuhjggg's picture
Upload 28 files
e83e49f verified
raw
history blame
1.17 kB
from flask import jsonify, send_file, request
from main import *
import torch
def summarize_text(text, output_path="output_summary.txt"):
if summarization_model is None or summarization_tokenizer is None:
return "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)
with open(output_path, "w") as file:
file.write(summary_text)
return output_path
def summarization_api():
data = request.get_json()
text = data.get('text')
if not text:
return jsonify({"error": "Text is required"}), 400
output_file = summarize_text(text)
if output_file == "Summarization model or tokenizer not initialized.":
return jsonify({"error": "Summarization failed"}), 500
return send_file(output_file, mimetype="text/plain", as_attachment=True, download_name="output_summary.txt")