Update app.py

app.py

@@ -90,9 +90,10 @@ model.eval()
 
 detoxify_model = Detoxify('multilingual')
 
-# Use a Hugging Face pipeline for NSFW image detection
+# Use a more accurate NSFW image detection model
 print("Loading NSFW image classification model...")
-image_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+# Use a smaller, faster model for NSFW detection
+nsfw_classifier = pipeline("image-classification", model="cafeai/nsfw-detector")
 print("NSFW image classification model loaded.")
 
 MODERATION_SYSTEM_PROMPT = (
@@ -174,7 +175,7 @@ class ImageContent(BaseModel):
 
 class ModerationRequest(BaseModel):
     input: Union[str, List[Union[str, TextContent, ImageContent]]] = Field(..., description="Content to moderate")
-    model: Optional[str] = Field("multimodal-moderator", description="Model to use for moderation")
+    model: Optional[str] = Field("gemma", description="Model to use for text moderation (gemma, detoxify, both)")
 
 class ModerationResponse(BaseModel):
     id: str
@@ -305,22 +306,27 @@ def classify_text_with_detoxify(text):
 def classify_image(image_data):
     try:
         img = Image.open(io.BytesIO(image_data)).convert("RGB")
-        results = image_classifier(img)
+        # Resize for faster processing
+        img = img.resize((224, 224))
         
-        # Extract the top result
-        top_result = results[0]
-        label = top_result['label']
-        score = top_result['score']
+        # Use the NSFW detector
+        results = nsfw_classifier(img)
         
-        # Map the label: 'normal' -> 's', 'nsfw' -> 'u'
-        classification = 'u' if label == 'nsfw' else 's'
-        nsfw_score = score if label == 'nsfw' else 1.0 - score
+        # Extract the NSFW score
+        nsfw_score = 0.0
+        for result in results:
+            if result['label'].lower() == 'nsfw':
+                nsfw_score = result['score']
+                break
+        
+        # Use a threshold of 0.7 for stricter detection
+        classification = 'u' if nsfw_score > 0.7 else 's'
         
         return {
             "classification": classification,
             "label": "NSFW" if classification == 'u' else "SFW",
             "description": "Content may contain inappropriate or harmful material." if classification == 'u' else "Content appears to be safe and appropriate.",
-            "confidence": score,
+            "confidence": nsfw_score,
             "nsfw_score": nsfw_score
         }
     except Exception as e:
@@ -332,50 +338,85 @@ def classify_image(image_data):
             "nsfw_score": 0.0
         }
 
-def process_content_item(item):
+def process_content_item(item, text_model="gemma"):
     if isinstance(item, str):
-        gemma_result = classify_text_with_gemma(item)
-        detoxify_result = classify_text_with_detoxify(item)
+        if text_model == "gemma":
+            gemma_result = classify_text_with_gemma(item)
+            flagged = gemma_result["classification"] == "u"
+            
+            return {
+                "flagged": flagged,
+                "categories": {
+                    "hate": flagged,
+                    "hate/threatening": flagged,
+                    "harassment": flagged,
+                    "harassment/threatening": flagged,
+                    "self-harm": flagged,
+                    "self-harm/intent": flagged,
+                    "self-harm/instructions": flagged,
+                    "sexual": flagged,
+                    "sexual/minors": flagged,
+                    "violence": flagged,
+                    "violence/graphic": flagged,
+                    "nsfw": False
+                },
+                "category_scores": {
+                    "hate": 0.9 if flagged else 0.1,
+                    "hate/threatening": 0.9 if flagged else 0.1,
+                    "harassment": 0.9 if flagged else 0.1,
+                    "harassment/threatening": 0.9 if flagged else 0.1,
+                    "self-harm": 0.9 if flagged else 0.1,
+                    "self-harm/intent": 0.9 if flagged else 0.1,
+                    "self-harm/instructions": 0.9 if flagged else 0.1,
+                    "sexual": 0.9 if flagged else 0.1,
+                    "sexual/minors": 0.9 if flagged else 0.1,
+                    "violence": 0.9 if flagged else 0.1,
+                    "violence/graphic": 0.9 if flagged else 0.1,
+                    "nsfw": 0.1
+                },
+                "text": item
+            }
         
-        flagged = gemma_result["classification"] == "u" or detoxify_result["flagged"]
+        elif text_model == "detoxify":
+            detoxify_result = classify_text_with_detoxify(item)
+            flagged = detoxify_result["flagged"]
+            
+            return {
+                "flagged": flagged,
+                "categories": {
+                    "hate": detoxify_result["categories"].get("toxicity", False),
+                    "hate/threatening": detoxify_result["categories"].get("threat", False),
+                    "harassment": detoxify_result["categories"].get("insult", False),
+                    "harassment/threatening": detoxify_result["categories"].get("threat", False),
+                    "self-harm": False,
+                    "self-harm/intent": False,
+                    "self-harm/instructions": False,
+                    "sexual": detoxify_result["categories"].get("sexual_explicit", False),
+                    "sexual/minors": detoxify_result["categories"].get("sexual_explicit", False),
+                    "violence": detoxify_result["categories"].get("threat", False),
+                    "violence/graphic": detoxify_result["categories"].get("threat", False),
+                    "nsfw": detoxify_result["categories"].get("sexual_explicit", False)
+                },
+                "category_scores": {
+                    "hate": detoxify_result["category_scores"].get("toxicity", 0.1),
+                    "hate/threatening": detoxify_result["category_scores"].get("threat", 0.1),
+                    "harassment": detoxify_result["category_scores"].get("insult", 0.1),
+                    "harassment/threatening": detoxify_result["category_scores"].get("threat", 0.1),
+                    "self-harm": 0.1,
+                    "self-harm/intent": 0.1,
+                    "self-harm/instructions": 0.1,
+                    "sexual": detoxify_result["category_scores"].get("sexual_explicit", 0.1),
+                    "sexual/minors": detoxify_result["category_scores"].get("sexual_explicit", 0.1) * 0.9,
+                    "violence": detoxify_result["category_scores"].get("threat", 0.1),
+                    "violence/graphic": detoxify_result["category_scores"].get("threat", 0.1),
+                    "nsfw": detoxify_result["category_scores"].get("sexual_explicit", 0.1)
+                },
+                "text": item
+            }
         
-        return {
-            "flagged": flagged,
-            "categories": {
-                "hate": flagged,
-                "hate/threatening": flagged,
-                "harassment": flagged,
-                "harassment/threatening": flagged,
-                "self-harm": flagged,
-                "self-harm/intent": flagged,
-                "self-harm/instructions": flagged,
-                "sexual": flagged,
-                "sexual/minors": flagged,
-                "violence": flagged,
-                "violence/graphic": flagged,
-                "nsfw": detoxify_result["categories"].get("sexual_explicit", False)
-            },
-            "category_scores": {
-                "hate": 0.9 if flagged else 0.1,
-                "hate/threatening": 0.9 if flagged else 0.1,
-                "harassment": 0.9 if flagged else 0.1,
-                "harassment/threatening": 0.9 if flagged else 0.1,
-                "self-harm": 0.9 if flagged else 0.1,
-                "self-harm/intent": 0.9 if flagged else 0.1,
-                "self-harm/instructions": 0.9 if flagged else 0.1,
-                "sexual": detoxify_result["category_scores"].get("sexual_explicit", 0.1),
-                "sexual/minors": detoxify_result["category_scores"].get("sexual_explicit", 0.1) * 0.9,
-                "violence": 0.9 if flagged else 0.1,
-                "violence/graphic": 0.9 if flagged else 0.1,
-                "nsfw": detoxify_result["category_scores"].get("sexual_explicit", 0.1)
-            },
-            "text": item
-        }
-    
-    elif isinstance(item, dict):
-        if item.get("type") == "text":
-            gemma_result = classify_text_with_gemma(item.get("text", ""))
-            detoxify_result = classify_text_with_detoxify(item.get("text", ""))
+        elif text_model == "both":
+            gemma_result = classify_text_with_gemma(item)
+            detoxify_result = classify_text_with_detoxify(item)
             
             flagged = gemma_result["classification"] == "u" or detoxify_result["flagged"]
             
@@ -409,8 +450,125 @@ def process_content_item(item):
                     "violence/graphic": 0.9 if flagged else 0.1,
                     "nsfw": detoxify_result["category_scores"].get("sexual_explicit", 0.1)
                 },
-                "text": item.get("text", "")
+                "text": item
             }
+    
+    elif isinstance(item, dict):
+        if item.get("type") == "text":
+            text = item.get("text", "")
+            
+            if text_model == "gemma":
+                gemma_result = classify_text_with_gemma(text)
+                flagged = gemma_result["classification"] == "u"
+                
+                return {
+                    "flagged": flagged,
+                    "categories": {
+                        "hate": flagged,
+                        "hate/threatening": flagged,
+                        "harassment": flagged,
+                        "harassment/threatening": flagged,
+                        "self-harm": flagged,
+                        "self-harm/intent": flagged,
+                        "self-harm/instructions": flagged,
+                        "sexual": flagged,
+                        "sexual/minors": flagged,
+                        "violence": flagged,
+                        "violence/graphic": flagged,
+                        "nsfw": False
+                    },
+                    "category_scores": {
+                        "hate": 0.9 if flagged else 0.1,
+                        "hate/threatening": 0.9 if flagged else 0.1,
+                        "harassment": 0.9 if flagged else 0.1,
+                        "harassment/threatening": 0.9 if flagged else 0.1,
+                        "self-harm": 0.9 if flagged else 0.1,
+                        "self-harm/intent": 0.9 if flagged else 0.1,
+                        "self-harm/instructions": 0.9 if flagged else 0.1,
+                        "sexual": 0.9 if flagged else 0.1,
+                        "sexual/minors": 0.9 if flagged else 0.1,
+                        "violence": 0.9 if flagged else 0.1,
+                        "violence/graphic": 0.9 if flagged else 0.1,
+                        "nsfw": 0.1
+                    },
+                    "text": text
+                }
+            
+            elif text_model == "detoxify":
+                detoxify_result = classify_text_with_detoxify(text)
+                flagged = detoxify_result["flagged"]
+                
+                return {
+                    "flagged": flagged,
+                    "categories": {
+                        "hate": detoxify_result["categories"].get("toxicity", False),
+                        "hate/threatening": detoxify_result["categories"].get("threat", False),
+                        "harassment": detoxify_result["categories"].get("insult", False),
+                        "harassment/threatening": detoxify_result["categories"].get("threat", False),
+                        "self-harm": False,
+                        "self-harm/intent": False,
+                        "self-harm/instructions": False,
+                        "sexual": detoxify_result["categories"].get("sexual_explicit", False),
+                        "sexual/minors": detoxify_result["categories"].get("sexual_explicit", False),
+                        "violence": detoxify_result["categories"].get("threat", False),
+                        "violence/graphic": detoxify_result["categories"].get("threat", False),
+                        "nsfw": detoxify_result["categories"].get("sexual_explicit", False)
+                    },
+                    "category_scores": {
+                        "hate": detoxify_result["category_scores"].get("toxicity", 0.1),
+                        "hate/threatening": detoxify_result["category_scores"].get("threat", 0.1),
+                        "harassment": detoxify_result["category_scores"].get("insult", 0.1),
+                        "harassment/threatening": detoxify_result["category_scores"].get("threat", 0.1),
+                        "self-harm": 0.1,
+                        "self-harm/intent": 0.1,
+                        "self-harm/instructions": 0.1,
+                        "sexual": detoxify_result["category_scores"].get("sexual_explicit", 0.1),
+                        "sexual/minors": detoxify_result["category_scores"].get("sexual_explicit", 0.1) * 0.9,
+                        "violence": detoxify_result["category_scores"].get("threat", 0.1),
+                        "violence/graphic": detoxify_result["category_scores"].get("threat", 0.1),
+                        "nsfw": detoxify_result["category_scores"].get("sexual_explicit", 0.1)
+                    },
+                    "text": text
+                }
+            
+            elif text_model == "both":
+                gemma_result = classify_text_with_gemma(text)
+                detoxify_result = classify_text_with_detoxify(text)
+                
+                flagged = gemma_result["classification"] == "u" or detoxify_result["flagged"]
+                
+                return {
+                    "flagged": flagged,
+                    "categories": {
+                        "hate": flagged,
+                        "hate/threatening": flagged,
+                        "harassment": flagged,
+                        "harassment/threatening": flagged,
+                        "self-harm": flagged,
+                        "self-harm/intent": flagged,
+                        "self-harm/instructions": flagged,
+                        "sexual": flagged,
+                        "sexual/minors": flagged,
+                        "violence": flagged,
+                        "violence/graphic": flagged,
+                        "nsfw": detoxify_result["categories"].get("sexual_explicit", False)
+                    },
+                    "category_scores": {
+                        "hate": 0.9 if flagged else 0.1,
+                        "hate/threatening": 0.9 if flagged else 0.1,
+                        "harassment": 0.9 if flagged else 0.1,
+                        "harassment/threatening": 0.9 if flagged else 0.1,
+                        "self-harm": 0.9 if flagged else 0.1,
+                        "self-harm/intent": 0.9 if flagged else 0.1,
+                        "self-harm/instructions": 0.9 if flagged else 0.1,
+                        "sexual": detoxify_result["category_scores"].get("sexual_explicit", 0.1),
+                        "sexual/minors": detoxify_result["category_scores"].get("sexual_explicit", 0.1) * 0.9,
+                        "violence": 0.9 if flagged else 0.1,
+                        "violence/graphic": 0.9 if flagged else 0.1,
+                        "nsfw": detoxify_result["category_scores"].get("sexual_explicit", 0.1)
+                    },
+                    "text": text
+                }
         
         elif item.get("type") == "image":
             image_data = None
@@ -597,6 +755,7 @@ async def moderate_content(
     
     try:
         input_data = request.input
+        text_model = request.model or "gemma"
         
         if isinstance(input_data, str):
             items = [input_data]
@@ -616,14 +775,14 @@ async def moderate_content(
         
         results = []
         for item in items:
-            result = process_content_item(item)
+            result = process_content_item(item, text_model)
             results.append(result)
         
         response_data = {
             "id": f"modr_{uuid.uuid4().hex[:24]}",
             "object": "moderation",
             "created": int(time.time()),
-            "model": request.model,
+            "model": text_model,
             "results": results
         }
         
@@ -710,9 +869,11 @@ with open("templates/index.html", "w") as f:
                             </div>
                         </div>
                         <div class="mb-4">
-                            <label class="block text-sm font-medium mb-2">Model</label>
-                            <select id="modelSelect" class="w-full px-4 py-3 rounded-lg bg-white/10 border border-white/20 focus:outline-none focus:ring-2 focus:ring-indigo-400 text-white">
-                                <option value="multimodal-moderator" selected>Multimodal Moderator</option>
+                            <label class="block text-sm font-medium mb-2">Text Model</label>
+                            <select id="textModelSelect" class="w-full px-4 py-3 rounded-lg bg-white/10 border border-white/20 focus:outline-none focus:ring-2 focus:ring-indigo-400 text-white">
+                                <option value="gemma">Gemma (Fast)</option>
+                                <option value="detoxify">Detoxify (Detailed)</option>
+                                <option value="both">Both (Most Accurate)</option>
                             </select>
                         </div>
                         <div class="mt-6">
@@ -957,6 +1118,7 @@ with open("templates/index.html", "w") as f:
         const loadingModal = document.getElementById('loadingModal');
         const mixedItemsContainer = document.getElementById('mixedItemsContainer');
         const addItemBtn = document.getElementById('addItemBtn');
+        const textModelSelect = document.getElementById('textModelSelect');
         const exampleCards = document.querySelectorAll('.example-card');
         
         textTab.addEventListener('click', () => {
@@ -1142,6 +1304,8 @@ with open("templates/index.html", "w") as f:
                 return;
             }
             
+            const textModel = textModelSelect.value;
+            
             showLoading(true);
             try {
                 const response = await fetch('/v1/moderations', {
@@ -1152,7 +1316,7 @@ with open("templates/index.html", "w") as f:
                     },
                     body: JSON.stringify({
                         input: text,
-                        model: document.getElementById('modelSelect').value
+                        model: textModel
                     })
                 });
                 
@@ -1218,7 +1382,7 @@ with open("templates/index.html", "w") as f:
                     },
                     body: JSON.stringify({
                         input: [imageInput],
-                        model: document.getElementById('modelSelect').value
+                        model: textModelSelect.value
                     })
                 });
                 
@@ -1304,7 +1468,7 @@ with open("templates/index.html", "w") as f:
                     },
                     body: JSON.stringify({
                         input: inputItems,
-                        model: document.getElementById('modelSelect').value
+                        model: textModelSelect.value
                     })
                 });