Update app.py

app.py

@@ -5,14 +5,13 @@ from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 # -------------------------
 # CONFIGURATION
 # -------------------------
-# The model you want to use (must have access from Hugging Face)
-MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
+# Switched to Qwen open-access model
+MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
 
 # -------------------------
 # TOKEN AUTHENTICATION
 # -------------------------
-# Your Hugging Face Access Token must be set in the HF Space as a Secret named "HF_TOKEN"
-# To do this, go to your Hugging Face Space > Settings > Secrets > Add "HF_TOKEN"
+# This model still requires your Hugging Face token (must be added as secret named "HF_TOKEN")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 if not HF_TOKEN:
@@ -22,40 +21,53 @@ if not HF_TOKEN:
 # LOAD TOKENIZER & MODEL
 # -------------------------
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, token=HF_TOKEN)
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True,
+        token=HF_TOKEN
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True,
+        device_map="auto",   # Automatically chooses GPU if available
+        token=HF_TOKEN
+    )
 except Exception as e:
     raise RuntimeError(f"🚨 Failed to load model: {e}")
 
 # -------------------------
 # CREATE PIPELINE
 # -------------------------
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=100,
-    do_sample=True,
-    temperature=0.7,
-)
+try:
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=300,
+        do_sample=True,
+        temperature=0.7,
+    )
+except Exception as e:
+    raise RuntimeError(f"🚨 Failed to initialize pipeline: {e}")
 
 # -------------------------
 # MAIN ASSISTANT FUNCTION
 # -------------------------
 def ai_assistant(command: str) -> str:
     """
-    Takes a natural language command and returns the assistant's response.
+    Interprets a user's natural language command and returns a response.
+    Uses instruction-style prompting.
     """
-    prompt = (
-        "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n"
-        f"{command}"
-        "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
-    )
-
+    prompt = f"User: {command}\nAssistant:"
+    
     try:
         output = pipe(prompt)[0]["generated_text"]
-        # Parse only the assistant response
-        response = output.split("<|eot_id|>")[0].split("<|end_header_id|>\n")[-1].strip()
+        # Remove the prompt portion to isolate the assistant's answer
+        if "Assistant:" in output:
+            response = output.split("Assistant:")[-1].strip()
+        else:
+            response = output.strip()
         return response
     except Exception as e:
         return f"⚠️ Error: {e}"
@@ -67,8 +79,8 @@ demo = gr.Interface(
     fn=ai_assistant,
     inputs=gr.Textbox(lines=2, placeholder="e.g. Open Chrome or Take a screenshot"),
     outputs="text",
-    title="🧠 LLaMA 3.1 AI Assistant",
-    description="Enter a command. The AI assistant will interpret and respond like a smart OS assistant.",
+    title="🧠 Qwen 2.5 AI Assistant",
+    description="Enter a command. The AI assistant will respond like a smart OS assistant.",
     allow_flagging="never"
 )