Update app/main.py

app/main.py

@@ -1,24 +1,17 @@
+import os
+os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")  # Prevents numba.cuda from initializing CUDA early
+
+import spaces  # Must come before any torch imports
+
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from glin_profanity import Filter
 import re
+import torch
 
-# 1. Toxicity classifier
+# Toxicity classifier — can stay at module level (CPU-based)
 toxicity_pipe = pipeline("text-classification", model="textdetox/twitter-xlmr-toxicity-classifier")
-
-# 2. Phi-3 Model - The FIX is here
-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
-    device_map="auto",            # Requires accelerate
-    torch_dtype="auto",           # Use the best dtype automatically
-)
-llm_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
 profanity_filter = Filter({"languages": ["english", "spanish", "french"], "replace_with": "***"})
 
 app = FastAPI(title="Multilingual Moderation API")
@@ -39,6 +32,44 @@ THREAT_KEYWORDS = [
 ]
 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"
+
+    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 = 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,
+    )
+    generated_text = tokenizer.decode(outputs[0][len(tokenizer(prompt)["input_ids"]):], skip_special_tokens=True).strip()
+    
+    lines = generated_text.split('\n')
+    rewritten = []
+    for line in lines:
+        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)
+    
+    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
@@ -57,28 +88,7 @@ async def moderate_text(request: TextRequest):
     
     rewritten = None
     if is_flagged:
-        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 = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        outputs = llm_pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
-        generated_text = outputs[0]['generated_text'][len(prompt):].strip()
-        
-        lines = generated_text.split('\n')
-        rewritten = []
-        for line in lines:
-            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)
-        
-        rewritten = list(dict.fromkeys(rewritten))[:5]
-        if not rewritten:
-            rewritten = ["[AI could not generate a rewrite, please revise manually]"]
+        rewritten = generate_suggestions(text)
 
     return ModerationResponse(
         original_text=text,