Update app.py

app.py

@@ -1,16 +1,14 @@
 import os
+import torch
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # -------------------------
 # CONFIGURATION
 # -------------------------
-MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"  # ✅ Text-only version
-
-# -------------------------
-# TOKEN AUTHENTICATION
-# -------------------------
+MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"  # small + fast for CPU
 HF_TOKEN = os.getenv("HF_TOKEN")
+
 if not HF_TOKEN:
     raise ValueError("🚫 Hugging Face token not found. Please add 'HF_TOKEN' in your Space secrets.")
 
@@ -18,17 +16,14 @@ if not HF_TOKEN:
 # LOAD TOKENIZER & MODEL
 # -------------------------
 try:
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_ID,
-        trust_remote_code=True,
-        token=HF_TOKEN
-    )
-
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
+        token=HF_TOKEN,
         trust_remote_code=True,
-        device_map="auto",
-        token=HF_TOKEN
+        device_map="cpu",
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
     )
 except Exception as e:
     raise RuntimeError(f"🚨 Failed to load model: {e}")
@@ -40,19 +35,24 @@ pipe = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_new_tokens=300,
+    max_new_tokens=100,
     do_sample=True,
-    temperature=0.7,
+    temperature=0.7
 )
 
 # -------------------------
 # MAIN ASSISTANT FUNCTION
 # -------------------------
 def ai_assistant(command: str) -> str:
-    prompt = f"User: {command}\nAssistant:"
+    prompt = f"<|user|>\n{command}\n<|assistant|>\n"
+
     try:
         output = pipe(prompt)[0]["generated_text"]
-        response = output.split("Assistant:")[-1].strip()
+        # Get assistant's part only
+        if "<|assistant|>" in output:
+            response = output.split("<|assistant|>")[-1].strip()
+        else:
+            response = output.strip()
         return response
     except Exception as e:
         return f"⚠️ Error: {e}"
@@ -64,9 +64,9 @@ demo = gr.Interface(
     fn=ai_assistant,
     inputs=gr.Textbox(lines=2, placeholder="e.g. Open Chrome or Take a screenshot"),
     outputs="text",
-    title="🧠 Qwen 2.5 AI Assistant",
-    description="Enter a command. The AI assistant will respond like a smart OS assistant.",
-    allow_flagging="never"
+    title="🧠 Qwen 1.5 AI Assistant",
+    description="Ask anything. This assistant uses the Qwen1.5-1.8B-Chat model on CPU.",
+    flagging_mode="never"
 )
 
 # -------------------------