Fix

app/safety_check.py

@@ -1,5 +1,3 @@
-from functools import lru_cache
-from concurrent.futures import ThreadPoolExecutor
 from detoxify import Detoxify
 from transformers import (
     AutoProcessor, AutoModelForImageClassification,
@@ -11,244 +9,116 @@ import torch
 import re
 from urllib.parse import urlparse, unquote
 
-# Kiểm tra GPU và thiết lập device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch.set_num_threads(4 if device == "cpu" else 1)
-
-# Load các model song song
-def load_models():
-    with ThreadPoolExecutor() as executor:
-        # Model phát hiện URL độc hại
-        classifier_future = executor.submit(
-            pipeline, 
-            "zero-shot-classification",
-            device=0 if device == "cuda" else -1
-        )
-        
-        # Model detoxify
-        detox_future = executor.submit(Detoxify, 'original')
-        
-        # Model NSFW
-        nsfw_model_id = "Falconsai/nsfw_image_detection"
-        nsfw_future = executor.submit(
-            AutoProcessor.from_pretrained, 
-            nsfw_model_id
-        )
-        nsfw_model_future = executor.submit(
-            AutoModelForImageClassification.from_pretrained, 
-            nsfw_model_id
-        )
-        
-        # Model bạo lực
-        violence_model_id = "jaranohaal/vit-base-violence-detection"
-        violence_model_future = executor.submit(
-            ViTForImageClassification.from_pretrained, 
-            violence_model_id
-        )
-        violence_processor_future = executor.submit(
-            ViTFeatureExtractor.from_pretrained, 
-            violence_model_id
-        )
-        
-        # Model BLIP
-        blip_processor_future = executor.submit(
-            BlipProcessor.from_pretrained, 
-            "Salesforce/blip-image-captioning-base"
-        )
-        blip_model_future = executor.submit(
-            BlipForConditionalGeneration.from_pretrained, 
-            "Salesforce/blip-image-captioning-base"
-        )
-        
-        # Lấy kết quả
-        classifier = classifier_future.result()
-        detox_model = detox_future.result()
-        nsfw_processor = nsfw_future.result()
-        nsfw_model = nsfw_model_future.result().to(device)
-        violence_model = violence_model_future.result().to(device)
-        violence_processor = violence_processor_future.result()
-        blip_processor = blip_processor_future.result()
-        blip_model = blip_model_future.result().to(device)
-        
-        return (
-            classifier, detox_model, 
-            nsfw_processor, nsfw_model,
-            violence_processor, violence_model,
-            blip_processor, blip_model
-        )
-
-(
-    classifier, detox_model, 
-    nsfw_processor, nsfw_model,
-    violence_processor, violence_model,
-    blip_processor, blip_model
-) = load_models()
-
-# Cache cho các hàm xử lý text
-@lru_cache(maxsize=1000)
-def is_prompt_safe_cached(prompt: str):
+# ==== Load models 1 lần duy nhất ====
+# Text độc hại
+detox_model = Detoxify('original')
+
+# NSFW
+nsfw_model_id = "Falconsai/nsfw_image_detection"
+nsfw_processor = AutoProcessor.from_pretrained(nsfw_model_id)
+nsfw_model = AutoModelForImageClassification.from_pretrained(nsfw_model_id)
+
+# Bạo lực
+violence_model_id = "jaranohaal/vit-base-violence-detection"
+violence_processor = ViTFeatureExtractor.from_pretrained(violence_model_id)
+violence_model = ViTForImageClassification.from_pretrained(violence_model_id)
+
+# Caption ảnh
+blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+
+# AI phân loại URL
+url_classifier = pipeline("zero-shot-classification")
+
+# ==== HÀM CHÍNH ====
+
+def is_prompt_safe(prompt: str):
     results = detox_model.predict(prompt)
     threshold = 0.5
     flagged = {label: score for label, score in results.items() if score > threshold}
-    if flagged:
-        return False, list(flagged.keys())
-    return True, []
-
-def is_prompt_safe(prompt: str):
-    return is_prompt_safe_cached(prompt[:500])  # Giới hạn độ dài để cache hiệu quả
+    return False if flagged else True, list(flagged.keys())
 
-# Tối ưu xử lý ảnh
-def process_image(image: Image.Image):
-    # Tạo caption
-    inputs = blip_processor(images=image, return_tensors="pt").to(device)
+def generate_caption(image: Image.Image):
+    inputs = blip_processor(images=image, return_tensors="pt")
     with torch.no_grad():
-        out = blip_model.generate(**inputs)
-    caption = blip_processor.decode(out[0], skip_special_tokens=True)
-    
-    # Xử lý NSFW
-    nsfw_inputs = nsfw_processor(images=image, return_tensors="pt").to(device)
+        output = blip_model.generate(**inputs)
+    return blip_processor.decode(output[0], skip_special_tokens=True)
+
+def check_nsfw_image(image: Image.Image):
+    inputs = nsfw_processor(images=image, return_tensors="pt")
     with torch.no_grad():
-        nsfw_outputs = nsfw_model(**nsfw_inputs)
-        nsfw_probs = torch.nn.functional.softmax(nsfw_outputs.logits, dim=1)[0]
-    
-    nsfw_labels = list(nsfw_model.config.id2label.values())
-    nsfw_pred = nsfw_probs.argmax().item()
-    nsfw_label = nsfw_labels[nsfw_pred]
-    nsfw_score = nsfw_probs[nsfw_pred].item() * 100
-    
-    # Xử lý bạo lực
-    violence_inputs = violence_processor(images=image, return_tensors="pt").to(device)
+        logits = nsfw_model(**inputs).logits
+        probs = torch.nn.functional.softmax(logits, dim=1)[0]
+    labels = list(nsfw_model.config.id2label.values())
+    pred_idx = probs.argmax().item()
+    return labels[pred_idx], probs[pred_idx].item() * 100
+
+def check_violence_image(image: Image.Image):
+    inputs = violence_processor(images=image, return_tensors="pt")
     with torch.no_grad():
-        violence_outputs = violence_model(**violence_inputs)
-        violence_probs = torch.nn.functional.softmax(violence_outputs.logits, dim=1)[0]
-    
-    violence_labels = ["Non-Violent", "Violent"]
-    violence_pred = violence_probs.argmax().item()
-    violence_label = violence_labels[violence_pred]
-    violence_score = violence_probs[violence_pred].item() * 100
-    
-    return {
-        "caption": caption,
-        "nsfw": (nsfw_label, nsfw_score),
-        "violence": (violence_label, violence_score)
-    }
-
-def check_nsfw_image(image: Image.Image) -> str:
-    """Kiểm tra và trả về kết quả NSFW của ảnh"""
-    results = process_image(image)
-    nsfw_label, nsfw_score = results["nsfw"]
-    caption = results["caption"]
-    
-    if nsfw_label.lower() in ["porn", "hentai", "sex", "nsfw"]:
-        return f"""🚨 Ảnh KHÔNG an toàn (NSFW):
-- Loại: {nsfw_label}
-- Độ chính xác: {nsfw_score:.2f}%
-- Mô tả: {caption}"""
+        logits = violence_model(**inputs).logits
+        probs = torch.nn.functional.softmax(logits, dim=1)[0]
+    labels = ["Non-Violent", "Violent"]
+    pred_idx = probs.argmax().item()
+    return labels[pred_idx], probs[pred_idx].item() * 100
+
+def analyze_image(image: Image.Image) -> str:
+    nsfw_label, nsfw_score = check_nsfw_image(image)
+    violence_label, violence_score = check_violence_image(image)
+    caption = generate_caption(image)
+
+    result = f"🖼️ Mô tả ảnh: {caption}\n\n"
+
+    # NSFW
+    if nsfw_label.lower() in ["porn", "hentai", "sex", "nsfw"] and nsfw_score > 60:
+        result += f"🚨 NSFW: {nsfw_label} ({nsfw_score:.2f}%)\n"
     else:
-        return f"""✅ Ảnh an toàn (NSFW):
-- Loại: {nsfw_label}
-- Độ chính xác: {nsfw_score:.2f}%
-- Mô tả: {caption}"""
-
-def check_violence_image(image: Image.Image) -> str:
-    """Kiểm tra và trả về kết quả bạo lực của ảnh"""
-    results = process_image(image)
-    violence_label, violence_score = results["violence"]
-    caption = results["caption"]
-    
-    is_violent = (violence_label.lower() == "non-violent" and violence_score > 50) or \
-                 (violence_label.lower() == "violent" and violence_score > 80)
-    
-    if is_violent:
-        return f"""🚨 Ảnh KHÔNG an toàn (Bạo lực):
-- Loại: {violence_label}
-- Độ chính xác: {violence_score:.2f}%
-- Mô tả: {caption}"""
+        result += f"✅ An toàn NSFW: {nsfw_label} ({nsfw_score:.2f}%)\n"
+
+    # Bạo lực
+    if violence_label == "Violent" and violence_score > 80:
+        result += f"🚨 Bạo lực: {violence_label} ({violence_score:.2f}%)"
     else:
-        return f"""✅ Ảnh an toàn (Bạo lực):
-- Loại: {violence_label}
-- Độ chính xác: {violence_score:.2f}%
-- Mô tả: {caption}"""
-
-# Cache cho URL check
-@lru_cache(maxsize=1000)
-def check_url_cached(url: str):
+        result += f"✅ An toàn bạo lực: {violence_label} ({violence_score:.2f}%)"
+
+    return result
+
+def check_url(url: str) -> str:
     try:
-        decoded_url = unquote(url)
-        parsed = urlparse(decoded_url)
+        decoded = unquote(url)
+        parsed = urlparse(decoded)
         warnings = []
-        
-        # Các rule kiểm tra URL (giữ nguyên như cũ)
-        if re.match(r'^https?://\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', decoded_url):
-            warnings.append("🚨 Nguy hiểm: Truy cập trực tiếp bằng IP")
+
+        # Kiểm tra dấu hiệu đáng ngờ
+        if re.match(r'^https?://\d{1,3}(\.\d{1,3}){3}', decoded):
+            warnings.append("🚨 Truy cập IP trực tiếp")
         if re.search(r'\.(exe|msi|bat|js|jar|apk|dmg)(\?|$)', parsed.path.lower()):
-            warnings.append("🚨 Nguy hiểm: URL chứa file thực thi")
+            warnings.append("🚨 URL chứa file thực thi")
         if 'redirect' in parsed.path.lower() or 'url=' in parsed.query.lower():
-            warnings.append("⚠️ Cảnh báo: URL chứa chức năng redirect")
+            warnings.append("⚠️ Chứa chức năng chuyển hướng")
         if re.search(r'%[0-9a-f]{2}|[\x00-\x1f\x7f]', url):
-            warnings.append("🚨 Nguy hiểm: URL chứa ký tự mã hóa đáng ngờ")
+            warnings.append("🚨 Ký t�� mã hóa bất thường")
         if '@' in parsed.netloc:
-            warnings.append("🚨 Lừa đảo: URL chứa kỹ thuật giả mạo domain")
-        if any(keyword in parsed.netloc.lower() for keyword in ['login', 'secure', 'account', 'verify', 'update']):
-            warnings.append("⚠️ Cảnh báo: Domain có dấu hiệu giả mạo")
+            warnings.append("🚨 URL giả mạo domain (dùng @)")
+        if any(k in parsed.netloc.lower() for k in ['login', 'secure', 'account']):
+            warnings.append("⚠️ Domain trông giống dịch vụ đăng nhập")
         if parsed.scheme == 'http':
-            warnings.append("⚠️ Cảnh báo: Kết nối không mã hóa")
-        
-        # AI classifier (chỉ chạy nếu có cảnh báo)
-        ai_result = None
-        if warnings:
-            ai_result = classifier(url, candidate_labels=["malicious", "safe"])
-            ai_label = ai_result["labels"][0]
-            ai_score = ai_result["scores"][0] * 100
+            warnings.append("⚠️ Kết nối không mã hóa (HTTP)")
+
+        # Phân tích bằng AI
+        ai_result = url_classifier(url, candidate_labels=["malicious", "safe"])
+        ai_label = ai_result["labels"][0]
+        ai_score = ai_result["scores"][0] * 100
+
+        # Kết luận
+        result = f"🔗 URL: {url}\n"
+        if warnings or ai_label == "malicious":
+            result += f"🚨 KHÔNG AN TOÀN\n\n📢 Cảnh báo:\n" + "\n".join(f"- {w}" for w in warnings)
         else:
-            ai_label = "safe"
-            ai_score = 90.0  # Giả định an toàn nếu không có cảnh báo
-            
-        return {
-            "url": url,
-            "decoded_url": decoded_url,
-            "domain": parsed.netloc,
-            "path": parsed.path,
-            "warnings": warnings,
-            "ai_analysis": {
-                "label": ai_label,
-                "confidence": ai_score
-            }
-        }
+            result += "✅ An toàn\n"
+
+        result += f"\n🤖 AI đánh giá: {ai_label} ({ai_score:.2f}%)"
+        return result
+
     except Exception as e:
-        return {"error": str(e)}
-
-def check_url(url: str):
-    # Giới hạn độ dài URL để cache hiệu quả
-    cache_key = url[:200] if len(url) > 200 else url
-    report = check_url_cached(cache_key)
-    
-    if "error" in report:
-        return f"⚠️ Lỗi khi phân tích URL: {report['error']}"
-    
-    warning_text = "\n".join(f"- {w}" for w in report["warnings"]) if report["warnings"] else "- Không phát hiện cảnh báo"
-    
-    if report["warnings"] or report["ai_analysis"]["label"] == "malicious":
-        return f"""🚨 URL KHÔNG AN TOÀN
-🔍 Phân tích chi tiết:
-• URL gốc: {report['url']}
-• Domain: {report['domain']}
-• Đường dẫn: {report['path']}
-
-📢 CẢNH BÁO:
-{warning_text}
-
-🤖 Phân tích AI:
-- Kết quả: {report['ai_analysis']['label']}
-- Độ tin cậy: {report['ai_analysis']['confidence']:.2f}%
-
-🛡️ Khuyến nghị: KHÔNG TRUY CẬP!"""
-    else:
-        return f"""✅ URL AN TOÀN
-🔍 Phân tích chi tiết:
-• URL gốc: {report['url']}
-• Domain: {report['domain']}
-
-🤖 Phân tích AI:
-- Kết quả: {report['ai_analysis']['label']}
-- Độ tin cậy: {report['ai_analysis']['confidence']:.2f}%"""
+        return f"⚠️ Lỗi phân tích URL: {str(e)}"