Spaces:
Sleeping
Sleeping
File size: 1,741 Bytes
da25abf 3944bf9 da25abf ad273e0 da25abf 4ac7937 da25abf 3944bf9 da25abf e70c72e da25abf e70c72e da25abf 4ac7937 da25abf |
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 47 48 49 50 51 52 53 54 55 56 57 |
from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import os
app = Flask(__name__)
# Dùng mô hình nhỏ hơn để tiết kiệm bộ nhớ
model_name = "VietAI/vit5-base-vietnews-summarization"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
summarizer = pipeline(
"summarization",
model=model,
tokenizer=tokenizer
)
@app.route("/")
def home():
return jsonify({
"message": "✅ AI Text Summarization API is running (Vietnamese).",
"usage": "POST /summarize với JSON: { 'text': '...', 'topic': '...' }"
})
@app.route("/summarize", methods=["POST"])
def summarize():
data = request.get_json(force=True)
if not data or "text" not in data:
return jsonify({"error": "Thiếu trường 'text'"}), 400
text = data["text"]
topic = data.get("topic", "").strip() # Nếu không có topic, mặc định trống
# Nếu có topic, ghép vào đầu văn bản để mô hình hiểu hướng tóm tắt
if topic:
text = f"Chủ đề: {topic}. Nội dung: {text}"
# Giới hạn số từ để tránh lỗi bộ nhớ
max_words = 1000
if len(text.split()) > max_words:
text = " ".join(text.split()[:max_words])
try:
# Tóm tắt văn bản
result = summarizer(text, max_length=150, min_length=40, do_sample=False)
summary = result[0]['summary_text']
return jsonify({"summary": summary})
except Exception as e:
return jsonify({"error": str(e)}), 500
if __name__ == "__main__":
port = int(os.environ.get("PORT", 7860))
app.run(host="0.0.0.0", port=port)
|