Update app.py

app.py

@@ -1,19 +1,79 @@
-
+import os
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
-model_id = "meta-llama/Llama-3.1-8B-Instruct"
+# -------------------------
+# CONFIGURATION
+# -------------------------
+# The model you want to use (must have access from Hugging Face)
+MODEL_ID = "meta-llama/Llama-3.1-8B-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"
+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.")
+
+# -------------------------
+# LOAD TOKENIZER & MODEL
+# -------------------------
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, token=HF_TOKEN)
+except Exception as e:
+    raise RuntimeError(f"🚨 Failed to load model: {e}")
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True, use_auth_token=True)
+# -------------------------
+# CREATE PIPELINE
+# -------------------------
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=100,
+    do_sample=True,
+    temperature=0.7,
+)
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# -------------------------
+# MAIN ASSISTANT FUNCTION
+# -------------------------
+def ai_assistant(command: str) -> str:
+    """
+    Takes a natural language command and returns the assistant's response.
+    """
+    prompt = (
+        "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n"
+        f"{command}"
+        "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
+    )
 
-def ai_assistant(command):
-    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n{command}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
-    result = pipe(prompt, max_new_tokens=100)[0]["generated_text"]
-    return result.split("<|eot_id|>")[0].split("<|end_header_id|>\n")[-1].strip()
+    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()
+        return response
+    except Exception as e:
+        return f"⚠️ Error: {e}"
 
-demo = gr.Interface(fn=ai_assistant, inputs="text", outputs="text", title="Llama 3.1 AI Assistant", description="Ask your assistant to do anything")
+# -------------------------
+# GRADIO UI
+# -------------------------
+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.",
+    allow_flagging="never"
+)
 
-demo.launch()
+# -------------------------
+# LAUNCH APP
+# -------------------------
+if __name__ == "__main__":
+    demo.launch()