Update app.py

app.py

@@ -1,30 +1,39 @@
-import os
-import gradio as gr
 import subprocess
+import os
 from llama_cpp import Llama
+import gradio as gr
 
+# 🔹 Get Hugging Face Token from environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Model download link (Modify if needed)
-MODEL_URL = "https://huggingface.co/TheBloke/mistral-7b.Q4_K_M.gguf"
+# 🔹 Model details
+MODEL_URL = "https://huggingface.co/TheBloke/Mistral-7B-GGUF/resolve/main/mistral-7b.Q4_K_M.gguf"
 MODEL_PATH = "./models/mistral-7b.Q4_K_M.gguf"
 
-# Create models directory if not exists
+# 🔹 Ensure the models directory exists
 os.makedirs("./models", exist_ok=True)
 
-# Auto-download model if not present
+# 🔹 Check if the model exists, else download it
 if not os.path.exists(MODEL_PATH):
-    print("Downloading Mistral-7B Q4 GGUF model...")
-    subprocess.run(["wget", MODEL_URL, "-O", MODEL_PATH], check=True)
+    print("🚀 Downloading Mistral-7B Q4 GGUF model...")
+    subprocess.run([
+        "wget", "--header", f"Authorization: Bearer {HF_TOKEN}",
+        MODEL_URL, "-O", MODEL_PATH
+    ], check=True)
+    print("✅ Download complete!")
 
-# Load GGUF model
-print("Loading model...")
+# 🔹 Load the model
+print("📥 Loading the model...")
 model = Llama(model_path=MODEL_PATH, n_ctx=4096, n_threads=8)
+print("✅ Model loaded successfully!")
 
-# Define function for chat
-def chat_with_ai(prompt):
-    response = model(prompt, max_tokens=512, stop=["</s>"])
+# 🔹 Define a function to interact with the model
+def chat_with_mistral(prompt):
+    response = model(prompt, max_tokens=512)
     return response["choices"][0]["text"]
 
-# Gradio UI
-iface = gr.Interface(fn=chat_with_ai, inputs="text", outputs="text", title="Mistral-7B GGUF Chatbot")
+# 🔹 Create a Gradio UI
+iface = gr.Interface(fn=chat_with_mistral, inputs="text", outputs="text", title="Mistral-7B Chatbot")
+
+# 🔹 Launch the app
 iface.launch()