Update app.py

app.py

@@ -7,8 +7,10 @@ import io
 import uuid
 import requests
 import numpy as np
+import asyncio
 from typing import List, Dict, Any, Optional, Union
 from fastapi import FastAPI, HTTPException, Depends, Request, File, UploadFile
+from fastapi.concurrency import run_in_threadpool
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
@@ -90,11 +92,10 @@ model.eval()
 
 detoxify_model = Detoxify('multilingual')
 
-# Load the NSFW image detection model and processor directly
 print("Loading NSFW image classification model...")
 nsfw_model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
 nsfw_processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
-nsfw_model.eval()  # Set to evaluation mode
+nsfw_model.eval()
 print("NSFW image classification model loaded.")
 
 MODERATION_SYSTEM_PROMPT = (
@@ -306,30 +307,24 @@ def classify_text_with_detoxify(text):
 
 def classify_image(image_data):
     try:
-        # Open and convert the image
         img = Image.open(io.BytesIO(image_data)).convert("RGB")
         
-        # Process the image with the NSFW model
         with torch.no_grad():
             inputs = nsfw_processor(images=img, return_tensors="pt")
-            # Move to the same device as the model
             inputs = {k: v.to(nsfw_model.device) for k, v in inputs.items()}
             
             outputs = nsfw_model(**inputs)
             logits = outputs.logits
             
-            # Get the predicted label
             predicted_label = logits.argmax(-1).item()
             label = nsfw_model.config.id2label[predicted_label]
             
-            # Get the confidence score
             confidence = torch.softmax(logits, dim=-1)[0][predicted_label].item()
             
-            # Convert to our classification system
             if label.lower() == "nsfw":
                 classification = "u"
                 nsfw_score = confidence
-            else:  # normal
+            else:
                 classification = "s"
                 nsfw_score = 1.0 - confidence
         
@@ -350,375 +345,124 @@ def classify_image(image_data):
             "nsfw_score": 0.0
         }
 
-def process_content_item(item, text_model="gemma"):
-    # Handle string input (simple text)
+def process_content_item(item: Union[str, Dict], text_model: str = "gemma") -> Dict:
     if isinstance(item, str):
+        item = {"type": "text", "text": item}
+
+    content_type = item.get("type")
+
+    if content_type == "text":
+        text = item.get("text", "")
         if text_model == "gemma":
-            gemma_result = classify_text_with_gemma(item)
+            gemma_result = classify_text_with_gemma(text)
             flagged = gemma_result["classification"] == "u"
-            
+            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,
+            }
             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
+                "categories": {k: (v > 0.5) for k, v in scores.items()},
+                "category_scores": scores,
+                "text": text,
             }
-        
         elif text_model == "detoxify":
-            detoxify_result = classify_text_with_detoxify(item)
-            flagged = detoxify_result["flagged"]
-            
+            d = classify_text_with_detoxify(text)
+            scores = {
+                "hate": d["category_scores"].get("toxicity", 0.1), "hate/threatening": d["category_scores"].get("threat", 0.1),
+                "harassment": d["category_scores"].get("insult", 0.1), "harassment/threatening": d["category_scores"].get("threat", 0.1),
+                "self-harm": 0.1, "self-harm/intent": 0.1, "self-harm/instructions": 0.1,
+                "sexual": d["category_scores"].get("sexual_explicit", 0.1), "sexual/minors": d["category_scores"].get("sexual_explicit", 0.1),
+                "violence": d["category_scores"].get("threat", 0.1), "violence/graphic": d["category_scores"].get("threat", 0.1),
+                "nsfw": d["category_scores"].get("sexual_explicit", 0.1),
+            }
             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
+                "flagged": d["flagged"],
+                "categories": {k: (v > 0.5) for k, v in scores.items()},
+                "category_scores": scores,
+                "text": text,
             }
-        
         elif text_model == "both":
-            gemma_result = classify_text_with_gemma(item)
-            detoxify_result = classify_text_with_detoxify(item)
-            
+            gemma_result = classify_text_with_gemma(text)
+            detoxify_result = classify_text_with_detoxify(text)
             flagged = gemma_result["classification"] == "u" or detoxify_result["flagged"]
-            
+            scores = {
+                "hate": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("toxicity", 0.1)),
+                "hate/threatening": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("threat", 0.1)),
+                "harassment": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("insult", 0.1)),
+                "harassment/threatening": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("threat", 0.1)),
+                "self-harm": 0.9 if gemma_result["classification"] == "u" else 0.1,
+                "self-harm/intent": 0.9 if gemma_result["classification"] == "u" else 0.1,
+                "self-harm/instructions": 0.9 if gemma_result["classification"] == "u" else 0.1,
+                "sexual": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("sexual_explicit", 0.1)),
+                "sexual/minors": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("sexual_explicit", 0.1)),
+                "violence": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("threat", 0.1)),
+                "violence/graphic": max(0.9 if gemma_result["classification"] == "u" else 0.1, detoxify_result["category_scores"].get("threat", 0.1)),
+                "nsfw": detoxify_result["category_scores"].get("sexual_explicit", 0.1),
+            }
             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
+                "categories": {k: (v > 0.5) for k, v in scores.items()},
+                "category_scores": scores,
+                "text": text,
             }
-    
-    # Handle dictionary input (structured content)
-    elif isinstance(item, dict):
-        content_type = item.get("type")
+
+    elif content_type == "image":
+        image_data = None
+        image_url = item.get("url")
+        image_base64 = item.get("base64")
         
-        if content_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
-                }
+        if image_url:
+            try:
+                response = requests.get(image_url, timeout=10)
+                response.raise_for_status()
+                image_data = response.content
+            except requests.RequestException as e:
+                print(f"Error fetching image from URL {image_url}: {e}")
+        elif image_base64:
+            try:
+                if image_base64.startswith("data:image"):
+                    image_base64 = image_base64.split(",")[1]
+                image_data = base64.b64decode(image_base64)
+            except Exception as e:
+                print(f"Error decoding base64 image: {e}")
         
-        elif content_type == "image":
-            image_data = None
-            image_url = item.get("url")
-            image_base64 = item.get("base64")
-            
-            # Get image data from URL
-            if image_url:
-                try:
-                    response = requests.get(image_url)
-                    if response.status_code == 200:
-                        image_data = response.content
-                    else:
-                        print(f"Failed to fetch image from URL: {image_url}, status code: {response.status_code}")
-                except Exception as e:
-                    print(f"Error fetching image from URL: {str(e)}")
-            
-            # Get image data from base64
-            elif image_base64:
-                try:
-                    if image_base64.startswith("data:image"):
-                        base64_data = image_base64.split(",")[1]
-                    else:
-                        base64_data = image_base64
-                    
-                    image_data = base64.b64decode(base64_data)
-                except Exception as e:
-                    print(f"Error decoding base64 image: {str(e)}")
-            
-            # Process the image if we have data
-            if image_data:
-                image_result = classify_image(image_data)
-                flagged = image_result["classification"] == "u"
-                
-                return {
-                    "flagged": flagged,
-                    "categories": {
-                        "hate": False,
-                        "hate/threatening": False,
-                        "harassment": False,
-                        "harassment/threatening": False,
-                        "self-harm": False,
-                        "self-harm/intent": False,
-                        "self-harm/instructions": False,
-                        "sexual": flagged,
-                        "sexual/minors": flagged,
-                        "violence": False,
-                        "violence/graphic": False,
-                        "nsfw": flagged
-                    },
-                    "category_scores": {
-                        "hate": 0.1,
-                        "hate/threatening": 0.1,
-                        "harassment": 0.1,
-                        "harassment/threatening": 0.1,
-                        "self-harm": 0.1,
-                        "self-harm/intent": 0.1,
-                        "self-harm/instructions": 0.1,
-                        "sexual": image_result["nsfw_score"],
-                        "sexual/minors": image_result["nsfw_score"] * 0.9,
-                        "violence": 0.1,
-                        "violence/graphic": 0.1,
-                        "nsfw": image_result["nsfw_score"]
-                    },
-                    "image_url": image_url,
-                    "image_base64": image_base64[:50] + "..." if image_base64 and len(image_base64) > 50 else image_base64
-                }
-            else:
-                # Return error if no image data
-                return {
-                    "flagged": False,
-                    "categories": {
-                        "hate": False,
-                        "hate/threatening": False,
-                        "harassment": False,
-                        "harassment/threatening": False,
-                        "self-harm": False,
-                        "self-harm/intent": False,
-                        "self-harm/instructions": False,
-                        "sexual": False,
-                        "sexual/minors": False,
-                        "violence": False,
-                        "violence/graphic": False,
-                        "nsfw": False
-                    },
-                    "category_scores": {
-                        "hate": 0.1,
-                        "hate/threatening": 0.1,
-                        "harassment": 0.1,
-                        "harassment/threatening": 0.1,
-                        "self-harm": 0.1,
-                        "self-harm/intent": 0.1,
-                        "self-harm/instructions": 0.1,
-                        "sexual": 0.1,
-                        "sexual/minors": 0.1,
-                        "violence": 0.1,
-                        "violence/graphic": 0.1,
-                        "nsfw": 0.1
-                    },
-                    "image_url": image_url,
-                    "image_base64": image_base64[:50] + "..." if image_base64 and len(image_base64) > 50 else image_base64
-                }
+        if image_data:
+            image_result = classify_image(image_data)
+            flagged = image_result["classification"] == "u"
+            nsfw_score = image_result.get("nsfw_score", 0.1)
+            scores = {
+                "hate": 0.1, "hate/threatening": 0.1,
+                "harassment": 0.1, "harassment/threatening": 0.1,
+                "self-harm": 0.1, "self-harm/intent": 0.1, "self-harm/instructions": 0.1,
+                "sexual": nsfw_score, "sexual/minors": nsfw_score,
+                "violence": 0.1, "violence/graphic": 0.1,
+                "nsfw": nsfw_score,
+            }
+            return {
+                "flagged": flagged,
+                "categories": {k: (v > 0.5) for k, v in scores.items()},
+                "category_scores": scores,
+                "image_url": image_url,
+                "image_base64": image_base64[:50] + "..." if isinstance(image_base64, str) and len(image_base64) > 50 else None,
+            }
     
-    # Default return for invalid items
+    default_scores = {
+        "hate": 0.1, "hate/threatening": 0.1, "harassment": 0.1, "harassment/threatening": 0.1,
+        "self-harm": 0.1, "self-harm/intent": 0.1, "self-harm/instructions": 0.1,
+        "sexual": 0.1, "sexual/minors": 0.1, "violence": 0.1, "violence/graphic": 0.1,
+        "nsfw": 0.1
+    }
     return {
         "flagged": False,
-        "categories": {
-            "hate": False,
-            "hate/threatening": False,
-            "harassment": False,
-            "harassment/threatening": False,
-            "self-harm": False,
-            "self-harm/intent": False,
-            "self-harm/instructions": False,
-            "sexual": False,
-            "sexual/minors": False,
-            "violence": False,
-            "violence/graphic": False,
-            "nsfw": False
-        },
-        "category_scores": {
-            "hate": 0.1,
-            "hate/threatening": 0.1,
-            "harassment": 0.1,
-            "harassment/threatening": 0.1,
-            "self-harm": 0.1,
-            "self-harm/intent": 0.1,
-            "self-harm/instructions": 0.1,
-            "sexual": 0.1,
-            "sexual/minors": 0.1,
-            "violence": 0.1,
-            "violence/graphic": 0.1,
-            "nsfw": 0.1
-        }
+        "categories": {k: False for k in default_scores},
+        "category_scores": default_scores,
+        "error": f"Invalid or unprocessable item: {item}"
     }
 
 def get_api_key(request: Request):
@@ -755,16 +499,12 @@ async def moderate_content(
         input_data = request.input
         text_model = request.model or "gemma"
         
-        # Normalize input to a list of items
         items = []
-        
         if isinstance(input_data, str):
-            # Single string input
-            items = [input_data]
+            items.append(input_data)
             total_tokens += count_tokens(input_data)
         elif isinstance(input_data, list):
-            # List of items
-            items = input_data
+            items.extend(input_data)
             for item in items:
                 if isinstance(item, str):
                     total_tokens += count_tokens(item)
@@ -775,19 +515,16 @@ async def moderate_content(
         
         if len(items) > 10:
             raise HTTPException(status_code=400, detail="Too many input items. Maximum 10 allowed.")
-        
-        # Process each item individually
-        results = []
-        for item in items:
-            result = process_content_item(item, text_model)
-            results.append(result)
+
+        tasks = [run_in_threadpool(process_content_item, item, text_model) for item in items]
+        results = await asyncio.gather(*tasks)
         
         response_data = {
             "id": f"modr_{uuid.uuid4().hex[:24]}",
             "object": "moderation",
             "created": int(time.time()),
             "model": text_model,
-            "results": results
+            "results": list(results)
         }
         
         track_request_metrics(start_time, total_tokens)
@@ -801,7 +538,7 @@ async def moderate_content(
 async def get_metrics(api_key: str = Depends(get_api_key)):
     return get_performance_metrics()
 
-with open("templates/index.html", "w") as f:
+with open("templates/index.html", "w", encoding='utf-8') as f:
     f.write("""<!DOCTYPE html>
 <html lang="en">
 <head>