Spaces:
Sleeping
Sleeping
File size: 1,378 Bytes
37a43bb 5af1c25 ee21003 37a43bb 5af1c25 ee21003 5af1c25 ee21003 37a43bb c981891 ee21003 37a43bb c981891 37a43bb |
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 44 45 46 |
from flask import Flask, request, jsonify
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import os, time
app = Flask(__name__)
# Load from local 'model/' folder
tokenizer = AutoTokenizer.from_pretrained("./model")
model = AutoModelForSeq2SeqLM.from_pretrained("./model")
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
# Warm up
summarizer("Warm-up input.", max_length=30, min_length=5, do_sample=False)
@app.route("/summarize", methods=["POST"])
def summarize():
try:
data = request.get_json()
text = data["text"]
max_len = int(data.get("max_length", 150)) # Default to 150 if not provided
min_len = int(data.get("min_length", 30)) # Default to 30 if not provided
summary_output = summarizer(
text,
max_length=max_len,
min_length=min_len,
do_sample=False
)
summary = summary_output[0]["summary_text"]
# Token count of the summary
tokenized = tokenizer(summary, return_tensors="pt")
token_count = len(tokenized["input_ids"][0])
return jsonify({
"summary": summary,
"token_count": token_count
})
except Exception as e:
return jsonify({"error": str(e)}), 400
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
|