Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import torch
-from beeper_model import BeeperRoseGPT, generate  # assumed modular split
+from beeper_model import BeeperRoseGPT, generate
 from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file as load_safetensors
 
 # ----------------------------
 # 🔧 Load Model and Tokenizer
@@ -20,36 +21,57 @@ config = {
     "repetition_penalty": 1.1,
     "presence_penalty": 0.6,
     "frequency_penalty": 0.0,
+    "resid_dropout": 0.1,  # Add these for model init
+    "dropout": 0.0,
+    "grad_checkpoint": False,
     "tokenizer_path": "beeper.tokenizer.json"
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load weights from Hugging Face repo if not available locally
+# Load weights from Hugging Face repo
 repo_id = "AbstractPhil/beeper-rose-tinystories-6l-512d-ctx512"
 model_file = hf_hub_download(repo_id=repo_id, filename="beeper_rose_final.safetensors")
 tokenizer_file = hf_hub_download(repo_id=repo_id, filename="tokenizer.json")
 
+# Initialize model
 infer = BeeperRoseGPT(config).to(device)
-infer.load_state_dict(torch.load(model_file, map_location=device))
+
+# Load safetensors properly
+state_dict = load_safetensors(model_file, device=str(device))
+infer.load_state_dict(state_dict)
 infer.eval()
+
+# Load tokenizer
 tok = Tokenizer.from_file(tokenizer_file)
 
 # ----------------------------
 # 💬 Gradio Chat Wrapper
 # ----------------------------
 def beeper_reply(message, history, temperature, top_k, top_p):
-    prompt = "\n".join([f"User: {h[0]}\nBeeper: {h[1]}" for h in history if h[0] and h[1]])
-    prompt += f"\nUser: {message}\nBeeper:"
-
-    out = generate(
+    # Build conversation context
+    prompt_parts = []
+    for h in history:
+        if h[0]:  # User message exists
+            prompt_parts.append(f"User: {h[0]}")
+        if h[1]:  # Assistant response exists
+            prompt_parts.append(f"Beeper: {h[1]}")
+    
+    # Add current message
+    prompt_parts.append(f"User: {message}")
+    prompt_parts.append("Beeper:")
+    
+    prompt = "\n".join(prompt_parts)
+    
+    # Generate response
+    response = generate(
         model=infer,
         tok=tok,
         cfg=config,
         prompt=prompt,
         max_new_tokens=128,
         temperature=temperature,
-        top_k=top_k,
+        top_k=int(top_k),
         top_p=top_p,
         repetition_penalty=config["repetition_penalty"],
         presence_penalty=config["presence_penalty"],
@@ -57,7 +79,12 @@ def beeper_reply(message, history, temperature, top_k, top_p):
         device=device,
         detokenize=True
     )
-    yield out
+    
+    # Clean up response - remove the prompt part if it's included
+    if response.startswith(prompt):
+        response = response[len(prompt):].strip()
+    
+    return response
 
 # ----------------------------
 # 🖼️ Interface
@@ -69,8 +96,16 @@ demo = gr.ChatInterface(
         gr.Slider(1, 100, value=40, step=1, label="Top-k"),
         gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
     ],
-    chatbot=gr.Chatbot(label="Hello I'm Beeper (Rose-based LLM)! Please be friendly I don't know very much yet!")
+    chatbot=gr.Chatbot(label="Chat with Beeper 🤖"),
+    title="Beeper - A Rose-based Tiny Language Model",
+    description="Hello! I'm Beeper, a small language model trained with love and care. Please be patient with me - I'm still learning! 💕",
+    examples=[
+        ["Hello Beeper! How are you today?"],
+        ["Can you tell me a story about a robot?"],
+        ["What do you like to do for fun?"],
+    ],
+    theme=gr.themes.Soft(),
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()