Update app.py

app.py

@@ -5,37 +5,33 @@ import torch
 # =========================
 # 1️⃣ Load model & tokenizer
 # =========================
-model_name = "Qwen/Qwen2.5-0.5B-Instruct"  # Change to 1.5B, 3B, 7B if needed
+model_name = "Qwen/Qwen2.5-0.5B-Instruct"
 
-# Use the fast tokenizer for speed
+# Fast tokenizer for speed
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 
-# Load model
+# Load model with correct dtype
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    device_map="auto",      # Automatically choose GPU if available
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    load_in_8bit=False       # Change to True if GPU memory is limited
+    device_map="auto",  # Uses GPU if available, else CPU
+    dtype=torch.float32  # CPU inference works better with float32
 )
 
-# Optional: PyTorch 2.x compile for speed
+# Optional PyTorch 2.x compile (speeds up CPU inference)
 if torch.__version__.startswith("2"):
     model = torch.compile(model)
 
 # =========================
 # 2️⃣ Hardcoded system prompt
 # =========================
-SYSTEM_PROMPT = "You are a friendly AI assistant named Magda that gives helpful and polite answers."
+SYSTEM_PROMPT = "You are a friendly AI assistant that gives helpful and polite answers."
 
 # =========================
-# 3️⃣ Chat function
+# 3️⃣ Optimized chat function
 # =========================
 def chat(user_prompt: str):
-    """
-    user_prompt: string message from user
-    """
     messages = [
-        {"role": "system", "content": SYSTEM_PROMPT},  # Hardcoded system
+        {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": user_prompt}
     ]
 
@@ -46,18 +42,18 @@ def chat(user_prompt: str):
         add_generation_prompt=True
     )
 
-    # Encode input
+    # Encode input once
     inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-    # Generate response (fast settings)
+    # Faster generation settings
     outputs = model.generate(
         **inputs,
         max_new_tokens=128,   # smaller = faster
         do_sample=False,      # deterministic = faster
-        num_beams=1           # no beam search = faster
+        num_beams=1           # no beam search
     )
 
-    # Decode output
+    # Decode only the first sequence
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response