Add 3

Dockerfile

@@ -1,21 +1,16 @@
-# Base image
 FROM python:3.10-slim
-
-# Set working directory
 WORKDIR /app
 
-# Copy requirements and install
+# Install dependencies
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy app code and model
+# Copy code & model
 COPY app.py .
 COPY model_save/ ./model_save/
 COPY best_model.pt .
 
-# Expose ports
 EXPOSE 7860
-EXPOSE 5000
 
-# Start Gradio UI + Flask API
-CMD bash -c "python app.py & gunicorn -w 4 -b 0.0.0.0:5000 app:app"
+# Run app
+CMD ["python", "app.py"]

app.py

@@ -1,27 +1,26 @@
 import gradio as gr
-from flask import Flask, request, jsonify
 import torch
 from transformers import BertTokenizer, BertForSequenceClassification
+from fastapi import FastAPI
+from pydantic import BaseModel
+import uvicorn
 import threading
 
 # ------------------- Device -------------------
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ------------------- Load model -------------------
+# ------------------- Load Model -------------------
 tokenizer = BertTokenizer.from_pretrained("./model_save")
 model = BertForSequenceClassification.from_pretrained("./model_save")
 model.load_state_dict(torch.load("best_model.pt", map_location=device))
 model.to(device)
 model.eval()
 
-# ------------------- Prediction function -------------------
-def predict_offensive(text):
+# ------------------- Prediction Function -------------------
+def predict_offensive(text: str):
     encoded = tokenizer(
-        text,
-        return_tensors="pt",
-        truncation=True,
-        padding="max_length",
-        max_length=128
+        text, return_tensors="pt",
+        truncation=True, padding="max_length", max_length=128
     )
     input_ids = encoded["input_ids"].to(device)
     attention_mask = encoded["attention_mask"].to(device)
@@ -32,36 +31,31 @@ def predict_offensive(text):
     pred = torch.argmax(logits, dim=1).item()
     return "Offensive" if pred == 1 else "Not Offensive"
 
-# ------------------- Flask API -------------------
-app = Flask(__name__)
-
-@app.route("/predict", methods=["POST"])
-def api_predict():
-    data = request.json
-    if not data or "text" not in data:
-        return jsonify({"error": "Missing 'text' field"}), 400
-
-    text = data["text"]
-    prediction = predict_offensive(text)
-    return jsonify({"prediction": prediction})
-
 # ------------------- Gradio UI -------------------
 iface = gr.Interface(
     fn=predict_offensive,
     inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
     outputs="text",
     title="Offensive Language Detector",
-    description="Enter a sentence and the model predicts if it contains offensive language."
 )
 
 def run_gradio():
     iface.launch(server_name="0.0.0.0", server_port=7860, share=False, prevent_thread_lock=True)
 
+# ------------------- FastAPI -------------------
+app = FastAPI(title="Offensive Language API")
+
+class TextItem(BaseModel):
+    text: str
+
+@app.post("/predict")
+def api_predict(item: TextItem):
+    return {"prediction": predict_offensive(item.text)}
+
 # ------------------- Main -------------------
 if __name__ == "__main__":
     # Start Gradio in a separate thread
     threading.Thread(target=run_gradio).start()
-    
-    # Flask API will be served by Gunicorn (HF Spaces sẽ build và chạy)
-    # gunicorn -w 4 -b 0.0.0.0:5000 app:app
-    print("Flask API ready. Use Gunicorn to serve for concurrent requests.")
+
+    # Run FastAPI (HF Spaces sẽ expose /docs automatically)
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,6 +1,7 @@
 torch
 transformers
-flask
-gunicorn
 gradio
+fastapi
+uvicorn
 numpy
+pydantic