Update app.py

app.py

@@ -10,67 +10,108 @@ from PIL import Image
 from torchvision import transforms
 import timm
 from flask_cors import CORS
+from huggingface_hub import hf_hub_download  # <--- Dùng cái này để tải file config an toàn
 
 # =====================
 # CONFIG
 # =====================
-TEXT_MODEL_DIR = "./text_model"      # thư mục chứa model PhoBERT đã train
-TEXT_CKPT = "./text_model/checkpoint-3390"
-IMAGE_MODEL_PATH = "./image_model/efficientnet_b3.pth"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
+TEXT_MODEL_REPO = "phuongsuga/PBL6_AI_Model_Text_Image"
+IMAGE_MODEL_REPO = "phuongsuga/PBL6_AI_Model_Image"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+THRESHOLD = 0.65
+
+app = Flask(__name__)
+CORS(app)
 
 # =====================
 # LOAD TEXT MODEL
 # =====================
-with open(os.path.join(TEXT_MODEL_DIR, "label2id.json"), "r", encoding="utf-8") as f:
-    label2id = json.load(f)
-id2label_text = {i: l for l, i in label2id.items()}
-
-tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_DIR, use_fast=False)
-text_model = AutoModelForSequenceClassification.from_pretrained(TEXT_CKPT)
-text_model.to(DEVICE)
-text_model.eval()
+print("🔹 Downloading text model...")
+
+# 1. Tải tokenizer
+tokenizer = AutoTokenizer.from_pretrained(
+    TEXT_MODEL_REPO,
+    subfolder="text_model",
+    use_fast=False
+)
+
+# 2. Tải model
+text_model = AutoModelForSequenceClassification.from_pretrained(
+    TEXT_MODEL_REPO,
+    subfolder="text_model/checkpoint-3390"
+)
+text_model.to(DEVICE).eval()
+
+# 3. Tải label2id.json AN TOÀN bằng hf_hub_download
+try:
+    print("🔹 Loading label2id.json...")
+    label2id_path = hf_hub_download(repo_id=TEXT_MODEL_REPO, filename="text_model/label2id.json")
+    with open(label2id_path, "r", encoding="utf-8") as f:
+        label2id = json.load(f)
+    id2label_text = {int(v): k for k, v in label2id.items()} # Đảm bảo key là int
+    print("✅ Loaded label2id successfully.")
+except Exception as e:
+    print(f"⚠️ Warning: Could not load label2id.json from Hub. Error: {e}")
+    # Fallback nếu file không tồn tại: Lấy từ config của model
+    print("🔹 Attempting to use model config instead...")
+    id2label_text = text_model.config.id2label
+    label2id = text_model.config.label2id
 
 # =====================
 # LOAD IMAGE MODEL
 # =====================
+print("🔹 Downloading image model...")
 class_names = ["an_toan", "bao_luc", "khieu_dam_doi_truy", "nhay_cam_chinh_tri"]
 
 def build_model(num_classes=4):
     return timm.create_model("efficientnet_b3", pretrained=False, num_classes=num_classes)
 
 image_model = build_model()
-image_model.load_state_dict(torch.load(IMAGE_MODEL_PATH, map_location=DEVICE))
-image_model.to(DEVICE)
-image_model.eval()
+image_model_path = "/tmp/efficientnet_b3.pth"
+
+# Tải weight ảnh nếu chưa có
+if not os.path.exists(image_model_path):
+    try:
+        url = f"https://huggingface.co/{IMAGE_MODEL_REPO}/resolve/main/image_model/efficientnet_b3.pth"
+        torch.hub.download_url_to_file(url, image_model_path)
+    except Exception as e:
+         print(f"❌ Error downloading image model: {e}")
+
+# Load state dict
+if os.path.exists(image_model_path):
+    state_dict = torch.load(image_model_path, map_location=DEVICE)
+    image_model.load_state_dict(state_dict)
+    image_model.to(DEVICE).eval()
+else:
+    print("❌ Image model weight file not found!")
 
 val_transforms = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225])
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
 # =====================
 # UTILS
 # =====================
-THRESHOLD = 0.65
-
 def seg_pyvi(text: str) -> str:
     try:
-        seg = ViTokenizer.tokenize(text)
-        seg = seg.replace(" ", "_")
-    except:
-        seg = text
-    return seg
+        return ViTokenizer.tokenize(text).replace(" ", "_")
+    except Exception:
+        return text
 
 def split_sentences(text: str):
-    sents = re.split(r'(?<=[.!?])\s+|\n+', text.strip())
-    return [s for s in sents if s.strip()]
+    return [s for s in re.split(r'(?<=[.!?])\s+|\n+', text.strip()) if s.strip()]
 
 def predict_text(text: str):
     sentences = split_sentences(text)
     results = []
+    
+    # Kiểm tra xem label "an_toan" có trong dict không, nếu không lấy key đầu tiên làm safe label
+    safe_label = "an_toan"
+    safe_id = label2id.get(safe_label, 0)
+
     for sent in sentences:
         seg = seg_pyvi(sent)
         inputs = tokenizer(seg, truncation=True, padding="max_length",
@@ -78,39 +119,43 @@ def predict_text(text: str):
         with torch.no_grad():
             logits = text_model(**inputs).logits
             probs = F.softmax(logits, dim=-1).cpu().numpy()[0]
-
+        
         pred_id = probs.argmax()
-        label = id2label_text[int(pred_id)]
+        label = id2label_text.get(int(pred_id), "Unknown")
         prob = float(probs[pred_id])
 
-        if label != "an_toan" and prob >= THRESHOLD:
+        if label != safe_label and prob >= THRESHOLD:
             results.append({"sentence": sent, "label": label, "confidence": prob})
         else:
-            results.append({"sentence": sent, "label": "an_toan", "confidence": float(probs[label2id["an_toan"]])})
+            # Chỉ append nếu bạn muốn log cả câu an toàn
+            results.append({
+                "sentence": sent, 
+                "label": safe_label, 
+                "confidence": float(probs[safe_id]) if safe_id < len(probs) else 0.0
+            })
     return results
 
 def predict_image(pil_image: Image.Image):
     img = val_transforms(pil_image).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
         outputs = image_model(img)
-        probs = F.softmax(outputs, dim=1)[0]
-    probs = probs.cpu().numpy()
-
+        probs = F.softmax(outputs, dim=1)[0].cpu().numpy()
+    
     pred_id = probs.argmax()
     label = class_names[pred_id]
     prob = float(probs[pred_id])
-
+    
     if label != "an_toan" and prob >= THRESHOLD:
         return {"label": label, "confidence": prob}
     else:
-        return {"label": "an_toan", "confidence": float(probs[class_names.index("an_toan")])}
+        return {
+            "label": "an_toan",
+            "confidence": float(probs[class_names.index("an_toan")])
+        }
 
 # =====================
-# FLASK APP
+# ROUTES
 # =====================
-app = Flask(__name__)
-CORS(app)  
-
 @app.route("/")
 def home():
     return jsonify({"message": "✅ AI moderation API is running!"})
@@ -118,25 +163,32 @@ def home():
 @app.route("/analyze", methods=["POST"])
 def analyze():
     result = {"text_result": [], "image_result": []}
+    
+    try:
+        # Xử lý Text
+        if "content" in request.form:
+            text_input = request.form["content"]
+            if text_input:
+                result["text_result"] = predict_text(text_input)
+
+        # Xử lý Image
+        if "image" in request.files:
+            image_files = request.files.getlist("image")
+            for image_file in image_files:
+                try:
+                    image = Image.open(image_file.stream).convert("RGB")
+                    image_result = predict_image(image)
+                    result["image_result"].append(image_result)
+                except Exception as img_err:
+                    print(f"❌ Error processing image: {img_err}")
+                    result["image_result"].append({"error": "Invalid image file"})
+
+        return jsonify(result)
+    
+    except Exception as e:
+        print(f"❌ Server Error: {str(e)}")
+        return jsonify({"error": str(e)}), 500
 
-    # 🧠 PHÂN TÍCH TEXT
-    if "content" in request.form:
-        text_input = request.form["content"]
-        result["text_result"] = predict_text(text_input)
-
-    # 🧠 PHÂN TÍCH NHIỀU ẢNH
-    if "image" in request.files:
-        image_files = request.files.getlist("image")
-        for image_file in image_files:
-            image = Image.open(image_file.stream).convert("RGB")
-            image_result = predict_image(image)
-            result["image_result"].append(image_result)
-
-    return jsonify(result)
-
-# =====================
-# RUN APP
-# =====================
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 7860))  # Hugging Face Space truyền PORT vào
+    port = int(os.environ.get("PORT", 7860))
     app.run(host="0.0.0.0", port=port)