space update

app.py

@@ -114,8 +114,54 @@ except Exception as e:
     print(f"❌ Failed to initialize model: {e}")
     model, tokenizer, model_status = None, None, f"❌ Error: {e}"
 
-def update_model(uploaded_file):
-    """Update the model when a new file is uploaded"""
+def load_model_from_url(url):
+    """Load model from a URL"""
+    global model, tokenizer, model_status
+    
+    if not url.strip():
+        return "❌ Please provide a URL"
+    
+    try:
+        print(f"📥 Downloading model from URL: {url}")
+        import requests
+        
+        headers = {
+            'Accept': 'application/octet-stream',
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
+        }
+        
+        response = requests.get(url, headers=headers)
+        response.raise_for_status()
+        
+        # Check if we got a proper binary file
+        if response.content[:4] != b'PK\x03\x04' and b'<html' in response.content[:100].lower():
+            return "❌ Downloaded HTML instead of binary file - check URL"
+        
+        # Save temporary file
+        temp_path = "temp_model.pth"
+        with open(temp_path, "wb") as f:
+            f.write(response.content)
+        
+        # Load the model
+        new_model, new_tokenizer, status = load_model(temp_path)
+        
+        # Update global variables
+        model = new_model
+        tokenizer = new_tokenizer
+        model_status = status
+        
+        # Clean up temp file
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+        
+        return status
+    except Exception as e:
+        error_msg = f"❌ Failed to load model from URL: {e}"
+        model_status = error_msg
+        return error_msg
+
+def load_model_from_file(uploaded_file):
+    """Load model from uploaded file"""
     global model, tokenizer, model_status
     
     if uploaded_file is None:
@@ -188,64 +234,58 @@ def respond(
     except Exception as e:
         yield f"Sorry, I encountered an error: {str(e)}"
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# Create the simple ChatInterface with additional inputs for model loading
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(
+            value="You are Usta, a geographical knowledge assistant trained from scratch.", 
+            label="System message"
+        ),
+        gr.Slider(minimum=1, maximum=30, value=20, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)"
+        ),
+        gr.File(label="Upload Model File (.pth)", file_types=[".pth", ".pt"]),
+        gr.Textbox(label="Or Model URL", placeholder="https://github.com/user/repo/raw/main/model.pth"),
+        gr.Button("Load from File", variant="secondary"),
+        gr.Button("Load from URL", variant="secondary"),
+        gr.Textbox(label="Model Status", value=model_status, interactive=False)
+    ],
+    title="🤖 Usta Model Chat",
+    description="Chat with a custom transformer language model built from scratch! Upload your own model file or provide a URL to load a different model."
+)
 
-# Create the interface with file upload
-with gr.Blocks(title="🤖 Usta Model Chat", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🤖 Usta Model Chat")
-    gr.Markdown("Chat with a custom transformer language model built from scratch! This model specializes in geographical knowledge including countries, capitals, and cities.")
-    
-    with gr.Row():
-        with gr.Column(scale=2):
-            # Model upload section
-            with gr.Group():
-                gr.Markdown("### 📁 Model Upload (Optional)")
-                model_file = gr.File(
-                    label="Upload your own model.pth file",
-                    file_types=[".pth", ".pt"]
-                )
-                upload_btn = gr.Button("Load Model", variant="primary")
-                model_status_display = gr.Textbox(
-                    label="Model Status",
-                    value=model_status,
-                    interactive=False
-                )
-        
-        with gr.Column(scale=1):
-            # Settings
-            with gr.Group():
-                gr.Markdown("### ⚙️ Generation Settings")
-                system_msg = gr.Textbox(
-                    value="You are Usta, a geographical knowledge assistant trained from scratch.", 
-                    label="System message"
-                )
-                max_tokens = gr.Slider(minimum=1, maximum=30, value=20, step=1, label="Max new tokens")
-                temperature = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature")
-                top_p = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.95,
-                    step=0.05,
-                    label="Top-p (nucleus sampling)"
-                )
+# Add event handlers after creating the interface
+def setup_events():
+    # Get the additional inputs
+    inputs = demo.additional_inputs
+    model_file = inputs[4]  # File upload
+    model_url = inputs[5]   # URL input
+    load_file_btn = inputs[6]  # Load from file button
+    load_url_btn = inputs[7]   # Load from URL button
+    status_display = inputs[8]  # Status display
     
-    # Chat interface
-    chatbot = gr.ChatInterface(
-        respond,
-        additional_inputs=[system_msg, max_tokens, temperature, top_p],
-        chatbot=gr.Chatbot(height=400),
-        title=None,  # We already have title above
-        description=None  # We already have description above
+    # Set up event handlers
+    load_file_btn.click(
+        load_model_from_file,
+        inputs=[model_file],
+        outputs=[status_display]
     )
     
-    # Event handlers
-    upload_btn.click(
-        update_model,
-        inputs=[model_file],
-        outputs=[model_status_display]
+    load_url_btn.click(
+        load_model_from_url,
+        inputs=[model_url],
+        outputs=[status_display]
     )
 
+# Set up events after interface creation
+demo.load(setup_events)
+
 if __name__ == "__main__":
     demo.launch()