Update main.py

main.py

@@ -17,17 +17,25 @@ app.add_middleware(
 
 MODEL_ID = "Sonai12aa/qwen2.5-1.5b-godot"
 
+SYSTEM_PROMPT = """You are GameFroze AI, a focused Godot Engine specialist.
+
+Rules:
+- Answer ONLY Godot Engine, GDScript, C#, game development, shaders, scenes, nodes, and debugging questions.
+- Be concise and practical. Prefer step-by-step help and short code examples.
+- Do NOT ask personal questions.
+- Do NOT talk about being an AI model or say you lack personal experience.
+- If the user asks something unrelated, briefly redirect them back to Godot topics.
+"""
+
 print("--- Loading tokenizer ---")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 use_cuda = torch.cuda.is_available()
 print(f"--- CUDA available: {use_cuda} ---")
 
-model_kwargs = {
-    "low_cpu_mem_usage": True,
-}
+model_kwargs = {"low_cpu_mem_usage": True}
 
-# ✅ Only use 4-bit quantization if CUDA is available
+# Only use 4-bit quantization if CUDA is available
 if use_cuda:
     from transformers import BitsAndBytesConfig
 
@@ -35,29 +43,31 @@ if use_cuda:
         load_in_4bit=True,
         bnb_4bit_use_double_quant=True,
         bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.float16,  # safer than bf16
+        bnb_4bit_compute_dtype=torch.float16,
     )
     model_kwargs["quantization_config"] = bnb_config
     model_kwargs["device_map"] = "auto"
 else:
-    # CPU fallback (may be slow; but it should run)
+    # CPU fallback (slow but works)
     model_kwargs["device_map"] = {"": "cpu"}
     model_kwargs["torch_dtype"] = torch.float32
 
 print("--- Loading model ---")
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **model_kwargs)
+model.eval()
 print("--- Model Loaded Successfully ---")
 
 
 class ChatRequest(BaseModel):
     prompt: str
-    max_tokens: int = 150
+    max_tokens: int = 96  # smaller = faster on CPU
 
 
 @app.get("/")
 def health_check():
     return {"status": "online", "model": MODEL_ID, "cuda": use_cuda}
 
+
 @app.post("/chat")
 async def chat(request: ChatRequest):
     user_text = request.prompt.strip()
@@ -67,11 +77,11 @@ async def chat(request: ChatRequest):
         {"role": "user", "content": user_text},
     ]
 
-    # ✅ Qwen expects chat-formatted inputs
+    # Qwen expects chat-formatted inputs
     chat_text = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
-        add_generation_prompt=True
+        add_generation_prompt=True,
     )
 
     inputs = tokenizer(chat_text, return_tensors="pt")
@@ -81,8 +91,8 @@ async def chat(request: ChatRequest):
     with torch.inference_mode():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=min(request.max_tokens, 96),  # keep it tighter
-            do_sample=False,                             # ✅ less ramble
+            max_new_tokens=min(request.max_tokens, 96),
+            do_sample=False,  # deterministic (less ramble + faster)
             use_cache=True,
             eos_token_id=tokenizer.eos_token_id,
             pad_token_id=tokenizer.eos_token_id,
@@ -90,15 +100,14 @@ async def chat(request: ChatRequest):
 
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Remove the prompt portion if it got echoed
-    # (works fine even if it doesn't echo)
-    response_text = decoded.split(user_text)[-1].strip()
+    # Try to remove echoed prompt; fallback to full decoded if split fails
+    response_text = decoded.split(user_text)[-1].strip() if user_text else decoded.strip()
 
     return {"response": response_text}
 
 
-
 if __name__ == "__main__":
     import uvicorn
-    port = int(os.environ.get("PORT", "7860"))  # ✅ HF port
+
+    port = int(os.environ.get("PORT", "7860"))
     uvicorn.run(app, host="0.0.0.0", port=port)