Update app.py

app.py

@@ -1,18 +1,24 @@
+import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-# Load model and tokenizer
+# ✅ Load token from Hugging Face secret
+HF_TOKEN = os.environ.get("key")
+
+# ✅ Model ID
 model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+# ✅ Load tokenizer and model securely
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
+    use_auth_token=HF_TOKEN,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     device_map="auto"
 )
 
-# Helper to build prompt
+# 🧠 Prompt Builder
 def build_prompt(user_input, history):
     prompt = "You are a pirate chatbot who always responds in pirate speak!\n"
     for user_msg, bot_reply in history:
@@ -20,7 +26,7 @@ def build_prompt(user_input, history):
     prompt += f"User: {user_input}\nPirate:"
     return prompt
 
-# Chat function
+# 💬 Chat Handler
 def chat(user_input, history):
     prompt = build_prompt(user_input, history)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -29,17 +35,16 @@ def chat(user_input, history):
         **inputs,
         max_new_tokens=256,
         do_sample=True,
-        top_p=0.9,
         temperature=0.8,
+        top_p=0.9,
         pad_token_id=tokenizer.eos_token_id
     )
 
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Extract only the last bot message
-    pirate_reply = decoded.split("Pirate:")[-1].strip()
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    pirate_reply = response.split("Pirate:")[-1].strip()
     return pirate_reply
 
-# Gradio UI
+# 🧱 Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 🏴‍☠️ Talk to the Pirate Bot!")
     chatbot = gr.Chatbot()