Update app/main.py

app/main.py

@@ -1,145 +1,42 @@
-import os
-os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
-
-import spaces  # must come before torch/transformers
-
-import re
-import torch
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 from glin_profanity import Filter
 
-# ---- CPU-only toxicity classifier (runs outside @spaces.GPU) ----
-toxicity_pipe = pipeline(
-    "text-classification",
-    model="textdetox/twitter-xlmr-toxicity-classifier",
-    device=-1,  # force CPU
-)
+# Initialize your models
+toxicity_pipe = pipeline("text-classification", model="textdetox/twitter-xlmr-toxicity-classifier")
+detox_pipe = pipeline("text2text-generation", model="s-nlp/mt0-xl-detox-orpo")
 profanity_filter = Filter({"languages": ["english", "spanish", "french"], "replace_with": "***"})
 
-# ---- Phi-3 loaded once on CPU; moved to GPU only inside the GPU function ----
-PHI_ID = "microsoft/Phi-3-mini-4k-instruct"
-phi_tokenizer = AutoTokenizer.from_pretrained(PHI_ID, trust_remote_code=True)
-phi_model = AutoModelForCausalLM.from_pretrained(
-    PHI_ID,
-    trust_remote_code=True,
-    attn_implementation="eager",
-    torch_dtype=torch.float16,
-)
-
 app = FastAPI(title="Multilingual Moderation API")
 
-
 class TextRequest(BaseModel):
     text: str
 
-
 class ModerationResponse(BaseModel):
     original_text: str
     detected_categories: dict
     is_flagged: bool
-    rewritten_text: list[str] | None
-
-
-THREAT_KEYWORDS = [
-    r"\bguns?\b", r"\bweapons?\b", r"\bbombs?\b", r"\bexplosives?\b",
-    r"\bkill\b", r"\bmurder\b", r"\bshoot\b", r"\bstab\b",
-    r"\bdie\b", r"\bdeath\b", r"\bhurt\b",
-]
-threat_pattern = re.compile("|".join(THREAT_KEYWORDS), re.IGNORECASE)
-
-
-@spaces.GPU
-def generate_suggestions(text: str) -> list[str]:
-    """Generate 5 polite rewrites on GPU."""
-    phi_model.to("cuda")
-
-    messages = [
-        {
-            "role": "system",
-            "content": (
-                "You are a helpful, polite assistant. Rewrite the user's toxic message "
-                "into exactly 5 different safe, polite, and non-threatening sentences. "
-                "Number them 1 to 5. Do not include any other text in your response."
-            ),
-        },
-        {"role": "user", "content": text},
-    ]
-
-    prompt = phi_tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    inputs = phi_tokenizer(prompt, return_tensors="pt").to("cuda")
-
-    with torch.no_grad():
-        outputs = phi_model.generate(
-            **inputs,
-            max_new_tokens=200,
-            do_sample=True,
-            temperature=0.7,
-        )
-
-    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
-    generated_text = phi_tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
-
-    rewritten = []
-    for line in generated_text.split("\n"):
-        line = line.strip()
-        if line and line[0].isdigit():
-            clean = re.sub(r"^\d+[\.\-\)]\s*", "", line).strip().strip("\"'")
-            if clean:
-                rewritten.append(clean)
-
-    rewritten = list(dict.fromkeys(rewritten))[:5]
-    return rewritten if rewritten else ["[AI could not generate a rewrite, please revise manually]"]
-
+    rewritten_text: str | None
 
 @app.post("/moderate", response_model=ModerationResponse)
 async def moderate_text(request: TextRequest):
     text = request.text
-
     profanity_result = profanity_filter.check_profanity(text)
-    # tox_result = toxicity_pipe(text)[0]
-    # raw_result = toxicity_pipe(text)
-    # tox_result = raw_result[0][0] if isinstance(raw_result[0], list) else raw_result[0]
-
-    # is_toxic_by_model = (
-    #     tox_result["label"].lower() in ["toxic", "label_1"] and tox_result["score"] > 0.5
-    # )
-    raw_result = toxicity_pipe(text)
-    # Handle both nested list [[{...}]] and flat list [{...}] outputs safely    
-    if isinstance(raw_result, list) and len(raw_result) > 0:
-        tox_result = raw_result[0][0] if isinstance(raw_result[0], list) else raw_result[0]    
-    else:       
-        tox_result = {"label": "neutral", "score": 0.0}
-        
-    is_toxic_by_model = (        
-        isinstance(tox_result, dict)        
-        and tox_result.get("label", "").lower() in ["toxic", "label_1"]
-        and tox_result.get("score", 0.0) > 0.5    
-    )
+    toxicity_scores = toxicity_pipe(text)[0]
+    
+    is_flagged = profanity_result["contains_profanity"] or toxicity_scores['score'] > 0.5
+    rewritten = None
+    if is_flagged:
+        rewritten = detox_pipe(text, max_length=100)[0]['generated_text']
     
-    is_threat_by_keyword = bool(threat_pattern.search(text))
-
-    is_flagged = bool(
-        profanity_result["contains_profanity"] or is_toxic_by_model or is_threat_by_keyword
-    )
-
-    rewritten = generate_suggestions(text) if is_flagged else None
-
     return ModerationResponse(
         original_text=text,
-        detected_categories={
-            "profanity": profanity_result["contains_profanity"],
-            "ai_toxicity": is_toxic_by_model,
-            "threat_keyword": is_threat_by_keyword,
-        },
+        detected_categories={"profanity": profanity_result["contains_profanity"], "toxicity_model": toxicity_scores},
         is_flagged=is_flagged,
-        rewritten_text=rewritten,
+        rewritten_text=rewritten
     )
 
-
 @app.get("/health")
 async def health_check():
     return {"status": "ok"}