Update app/main.py

app/main.py

@@ -1,106 +1,130 @@
 import os
-os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")  # Prevents numba.cuda from initializing CUDA early
+os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
 
-import spaces  # Must come before any torch imports
+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 glin_profanity import Filter
-import re
-import torch
 
-# Toxicity classifier — can stay at module level (CPU-based)
-toxicity_pipe = pipeline("text-classification", model="textdetox/twitter-xlmr-toxicity-classifier")
+# ---- CPU-only toxicity classifier (runs outside @spaces.GPU) ----
+toxicity_pipe = pipeline(
+    "text-classification",
+    model="textdetox/twitter-xlmr-toxicity-classifier",
+    device=-1,  # force CPU
+)
 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"
+    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]:
-    """Load Phi-3 and generate rewritten sentences on GPU."""
-    model_id = "microsoft/Phi-3-mini-4k-instruct"
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        trust_remote_code=True,
-        attn_implementation="eager",  # Solves the window_size error
-    ).to("cuda")  # Use .to("cuda"), NOT device_map="auto"
+    """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}
+        {
+            "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 = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    outputs = model.generate(
-        **tokenizer(prompt, return_tensors="pt").to("cuda"),
-        max_new_tokens=200,
-        do_sample=True,
-        temperature=0.7,
+
+    prompt = phi_tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
     )
-    generated_text = tokenizer.decode(outputs[0][len(tokenizer(prompt)["input_ids"]):], skip_special_tokens=True).strip()
-    
-    lines = generated_text.split('\n')
+    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 lines:
+    for line in generated_text.split("\n"):
         line = line.strip()
         if line and line[0].isdigit():
-            clean_line = re.sub(r'^\d+[\.\-\)]\s*', '', line).strip()
-            clean_line = clean_line.strip('"\'')
-            if clean_line:
-                rewritten.append(clean_line)
-    
+            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]"]
 
+
 @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]
-    
-    is_toxic_by_model = (tox_result['label'].lower() in ['toxic', 'label_1']) and (tox_result['score'] > 0.5)
+
+    is_toxic_by_model = (
+        tox_result["label"].lower() in ["toxic", "label_1"] and tox_result["score"] > 0.5
+    )
     is_threat_by_keyword = bool(threat_pattern.search(text))
-    
-    is_flagged = (
-        profanity_result["contains_profanity"] or 
-        is_toxic_by_model or 
-        is_threat_by_keyword
+
+    is_flagged = bool(
+        profanity_result["contains_profanity"] or is_toxic_by_model or is_threat_by_keyword
     )
-    
-    rewritten = None
-    if is_flagged:
-        rewritten = generate_suggestions(text)
+
+    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
+            "threat_keyword": is_threat_by_keyword,
         },
         is_flagged=is_flagged,
-        rewritten_text=rewritten
+        rewritten_text=rewritten,
     )
 
+
 @app.get("/health")
 async def health_check():
     return {"status": "ok"}