Create app.py

app.py

@@ -1,76 +1,20 @@
-import os
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from model_loader import load_model
+from assistant import get_assistant_response
 
-# -------------------------
-# CONFIGURATION
-# -------------------------
-MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"  # small + fast for CPU
-HF_TOKEN = os.getenv("HF_TOKEN")
+pipe = load_model()
 
-if not HF_TOKEN:
-    raise ValueError("🚫 Hugging Face token not found. Please add 'HF_TOKEN' in your Space secrets.")
-
-# -------------------------
-# LOAD TOKENIZER & MODEL
-# -------------------------
-try:
-    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="cpu",
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True
-    )
-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
-)
-
-# -------------------------
-# MAIN ASSISTANT FUNCTION
-# -------------------------
-def ai_assistant(command: str) -> str:
-    prompt = f"<|user|>\n{command}\n<|assistant|>\n"
-
-    try:
-        output = pipe(prompt)[0]["generated_text"]
-        # 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}"
-
-# -------------------------
-# GRADIO UI
-# -------------------------
 demo = gr.Interface(
-    fn=ai_assistant,
-    inputs=gr.Textbox(lines=2, placeholder="e.g. Open Chrome or Take a screenshot"),
+    fn=lambda command, execute: get_assistant_response(pipe, command, execute),
+    inputs=[
+        gr.Textbox(lines=2, placeholder="e.g. Open Chrome or Take a screenshot"),
+        gr.Checkbox(label="🛠️ Execute command (if possible)")
+    ],
     outputs="text",
-    title="🧠 Qwen 1.5 AI Assistant",
-    description="Ask anything. This assistant uses the Qwen1.5-1.8B-Chat model on CPU.",
+    title="🧠 Smart AI Assistant",
+    description="Try commands like 'Open Chrome', 'Take a screenshot', 'Create file', or 'What is AI?'",
     flagging_mode="never"
 )
 
-# -------------------------
-# LAUNCH APP
-# -------------------------
 if __name__ == "__main__":
     demo.launch()