Update app.py

app.py

@@ -6,8 +6,22 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file as load_safetensors
 
 # ----------------------------
-# 🔧 Load Model and Tokenizer
+# 🔧 Model versions configuration
 # ----------------------------
+MODEL_VERSIONS = {
+    "Beeper v1 (Original)": {
+        "repo_id": "AbstractPhil/beeper-rose-tinystories-6l-512d-ctx512",
+        "model_file": "beeper_rose_final.safetensors",
+        "description": "Original Beeper trained on TinyStories"
+    },
+    "Beeper v2 (Extended)": {
+        "repo_id": "AbstractPhil/beeper-rose-v2",
+        "model_file": "beeper_final.safetensors",
+        "description": "Beeper v2 with extended training (~15 epochs)"
+    }
+}
+
+# Base configuration
 config = {
     "context": 512,
     "vocab_size": 8192,
@@ -21,7 +35,7 @@ config = {
     "repetition_penalty": 1.1,
     "presence_penalty": 0.6,
     "frequency_penalty": 0.0,
-    "resid_dropout": 0.1,  # Add these for model init
+    "resid_dropout": 0.1,
     "dropout": 0.0,
     "grad_checkpoint": False,
     "tokenizer_path": "beeper.tokenizer.json"
@@ -29,26 +43,68 @@ config = {
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 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)
+# Global model and tokenizer variables
+infer = None
+tok = None
+current_version = None
 
-# Load safetensors properly
-state_dict = load_safetensors(model_file, device=str(device))
-infer.load_state_dict(state_dict)
-infer.eval()
+def load_model_version(version_name):
+    """Load the selected model version"""
+    global infer, tok, current_version
+    
+    if current_version == version_name and infer is not None:
+        return f"Already loaded: {version_name}"
+    
+    version_info = MODEL_VERSIONS[version_name]
+    
+    try:
+        # Download model and tokenizer files
+        model_file = hf_hub_download(
+            repo_id=version_info["repo_id"], 
+            filename=version_info["model_file"]
+        )
+        tokenizer_file = hf_hub_download(
+            repo_id=version_info["repo_id"], 
+            filename="tokenizer.json"
+        )
+        
+        # Initialize model
+        infer = BeeperRoseGPT(config).to(device)
+        
+        # Load safetensors
+        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)
+        
+        current_version = version_name
+        return f"Successfully loaded: {version_name}"
+    
+    except Exception as e:
+        return f"Error loading {version_name}: {str(e)}"
 
-# Load tokenizer
-tok = Tokenizer.from_file(tokenizer_file)
+# Load default model on startup
+load_status = load_model_version("Beeper v1 (Original)")
+print(load_status)
 
 # ----------------------------
 # 💬 Gradio Chat Wrapper
 # ----------------------------
-def beeper_reply(message, history, temperature=None, top_k=None, top_p=None):
+def beeper_reply(message, history, model_version, temperature=None, top_k=None, top_p=None):
+    global infer, tok, current_version
+    
+    # Load model if version changed
+    if model_version != current_version:
+        status = load_model_version(model_version)
+        if "Error" in status:
+            return f"⚠️ {status}"
+    
+    # Check if model is loaded
+    if infer is None or tok is None:
+        return "⚠️ Model not loaded. Please select a version and try again."
+    
     # Use defaults if not provided (for examples caching)
     if temperature is None:
         temperature = 0.9
@@ -98,24 +154,81 @@ def beeper_reply(message, history, temperature=None, top_k=None, top_p=None):
 # ----------------------------
 # 🖼️ Interface
 # ----------------------------
-demo = gr.ChatInterface(
-    beeper_reply,
-    additional_inputs=[
-        gr.Slider(0.1, 1.5, value=0.9, step=0.1, label="Temperature"),
-        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="Chat with Beeper 🤖", type="messages"),
-    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(),
-    cache_examples=False,  # Disable caching to avoid the startup issue
-)
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🤖 Beeper - A Rose-based Tiny Language Model
+        Hello! I'm Beeper, a small language model trained with love and care. Please be patient with me - I'm still learning! 💕
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            model_dropdown = gr.Dropdown(
+                choices=list(MODEL_VERSIONS.keys()),
+                value="Beeper v1 (Original)",
+                label="Select Beeper Version",
+                info="Choose which version of Beeper to chat with"
+            )
+        with gr.Column(scale=7):
+            version_info = gr.Markdown("**Current:** Beeper v1 - Original training on TinyStories")
+    
+    # Update version info when dropdown changes
+    def update_version_info(version_name):
+        info = MODEL_VERSIONS[version_name]["description"]
+        return f"**Current:** {info}"
+    
+    model_dropdown.change(
+        fn=update_version_info,
+        inputs=[model_dropdown],
+        outputs=[version_info]
+    )
+    
+    # Chat interface
+    chatbot = gr.Chatbot(label="Chat with Beeper", type="messages", height=400)
+    msg = gr.Textbox(label="Message", placeholder="Type your message here...")
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            temperature_slider = gr.Slider(0.1, 1.5, value=0.9, step=0.1, label="Temperature")
+        with gr.Column(scale=2):
+            top_k_slider = gr.Slider(1, 100, value=40, step=1, label="Top-k")
+        with gr.Column(scale=2):
+            top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+    
+    with gr.Row():
+        submit = gr.Button("Send", variant="primary")
+        clear = gr.Button("Clear")
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            ["Hello Beeper! How are you today?"],
+            ["Can you tell me a story about a robot?"],
+            ["What do you like to do for fun?"],
+            ["What makes you happy?"],
+            ["Tell me about your dreams"],
+        ],
+        inputs=msg
+    )
+    
+    # Handle chat
+    def respond(message, chat_history, model_version, temperature, top_k, top_p):
+        response = beeper_reply(message, chat_history, model_version, temperature, top_k, top_p)
+        chat_history.append([message, response])
+        return "", chat_history
+    
+    msg.submit(
+        respond, 
+        [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider], 
+        [msg, chatbot]
+    )
+    submit.click(
+        respond, 
+        [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider], 
+        [msg, chatbot]
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
     demo.launch()