Update app.py

app.py

@@ -8,7 +8,7 @@ import uvicorn
 # -------------------------------
 # SETTINGS
 # -------------------------------
-MODEL_NAME = "TheDrummer/Gemmasutra-Mini-2B-v1"  # get out lol
+MODEL_NAME = "TheDrummer/Gemmasutra-Mini-2B-v1"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_TOKENS = 512  # max tokens for a response
 
@@ -56,6 +56,9 @@ async def ask_ai(request: Request):
     data = await request.json()
     user_prompt = data.get("prompt", "")
     
+    if not user_prompt.strip():
+        return {"reply": "No prompt provided."}
+    
     # Combine system prompt + user input
     full_prompt = SYSTEM_PROMPT + "\nUser: " + user_prompt + "\nAI:"
     
@@ -64,28 +67,33 @@ async def ask_ai(request: Request):
     
     # Generate response
     outputs = model.generate(
-    input_ids,
-    max_new_tokens=300,
-    do_sample=False,
-    temperature=0.0,
-    top_p=1.0,
-    repetition_penalty=1.1,
-    eos_token_id=tokenizer.eos_token_id
-)
-
+        **inputs,                           # ← fixed: was using undefined 'input_ids'
+        max_new_tokens=300,
+        do_sample=False,
+        temperature=0.0,
+        top_p=1.0,
+        repetition_penalty=1.1,
+        eos_token_id=tokenizer.eos_token_id
+    )
     
     # Decode output
-    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    reply = reply.split("User:")[0].split("AI:")[0].strip()
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the new part after the prompt (most reliable method)
+    if generated_text.startswith(full_prompt):
+        reply = generated_text[len(full_prompt):].strip()
+    else:
+        # fallback - try to cut after last "AI:"
+        if "AI:" in generated_text:
+            reply = generated_text.split("AI:")[-1].strip()
+        else:
+            reply = generated_text.strip()
     
-    # Remove the system prompt from reply
-    reply = reply.replace(full_prompt, "").strip()
     return {"reply": reply}
 
 
-
 # -------------------------------
 # RUN SERVER
 # -------------------------------
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)