switch to fastapi

Dockerfile

@@ -18,5 +18,4 @@ COPY app.py .
 # Expose port mặc định HFS (7860)
 EXPOSE 7860
 
-# Run Flask
-CMD ["python", "app.py"]
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,46 +1,40 @@
-import os
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Khai báo thư mục cache an toàn cho Hugging Face
-os.environ["HF_HOME"] = "/app/cache"
-os.environ["TRANSFORMERS_CACHE"] = "/app/cache/transformers"
+import torch
 
-app = Flask(__name__)
+app = FastAPI()
 
+# Load model
 model_name = "VietAI/vit5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-@app.route("/summarize", methods=["POST"])
-def summarize():
-    data = request.get_json()
-    text = data.get("text", "").strip()
-
-    if not text:
-        return jsonify({"error": "Missing 'text' field"}), 400
-
-    # ✅ Rất quan trọng: Thêm tiền tố 'summarize:'
-    prompt = f"summarize: {text}"
-    inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
-
-    summary_ids = model.generate(
-        inputs,
-        max_length=100,
-        min_length=10,
-        num_beams=4,
-        no_repeat_ngram_size=3,
-        repetition_penalty=2.5,
-        length_penalty=1.0,
-        early_stopping=True
-    )
-
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return jsonify({"summary": summary})
-
-@app.route("/", methods=["GET"])
-def index():
-    return "✅ ViT5 summarization API is running."
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+# Input format
+class TextInput(BaseModel):
+    text: str
+
+@app.get("/")
+def read_root():
+    return {"message": "ViT5 summarization API is running!"}
+
+@app.post("/summarize")
+def summarize(input: TextInput):
+    try:
+        input_text = f"summarize: {input.text}"
+        inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
+        summary_ids = model.generate(
+            inputs,
+            max_length=128,
+            min_length=20,
+            num_beams=4,
+            no_repeat_ngram_size=3,
+            repetition_penalty=2.5,
+            length_penalty=1.0,
+            early_stopping=True
+        )
+        output = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+        return {"summary": output}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

requirements.txt

@@ -1,4 +1,4 @@
-flask
-transformers
+transformers==4.41.2
 torch
-sentencepiece
+fastapi
+uvicorn