space update

app.py

@@ -8,7 +8,7 @@ from v1.usta_tokenizer import UstaTokenizer
 
 
 # Load the model and tokenizer
-def load_model():
+def load_model(custom_model_path=None):
     try:
         u_tokenizer = UstaTokenizer("v1/tokenizer.json")
         print("✅ Tokenizer loaded successfully! vocab size:", len(u_tokenizer.vocab))
@@ -29,66 +29,96 @@ def load_model():
             num_layers=num_layers
         )        
         
-        # Load the trained weights if available
-        model_path = "v1/u_model.pth"
-
-        if not os.path.exists(model_path):
-            print("❌ Model file not found at", model_path)
-            # Download the model file from GitHub
-            try:
-                print("📥 Downloading model weights from GitHub...")
-                import requests
-                url = "https://github.com/malibayram/llm-from-scratch/raw/main/u_model_4000.pth"
-                
-                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()  # Raise an exception for bad status codes
-                
-                # Check if we got a proper binary file (PyTorch files start with specific bytes)
-                if response.content[:4] != b'PK\x03\x04' and b'<html' in response.content[:100].lower():
-                    raise Exception("Downloaded HTML instead of binary file - check URL")
-                
-                print(f"📦 Downloaded {len(response.content)} bytes")
-                
-                # Create v1 directory if it doesn't exist
-                os.makedirs("v1", exist_ok=True)
-                
-                # Save the model weights to the local file system
-                with open(model_path, "wb") as f:
-                    f.write(response.content)
-                print("✅ Model weights saved successfully!")
-            except Exception as e:
-                print(f"❌ Failed to download model weights: {e}")
-                print("Using random initialization.")
+        # Determine which model file to use
+        if custom_model_path and os.path.exists(custom_model_path):
+            model_path = custom_model_path
+            print(f"🎯 Using uploaded model: {model_path}")
+        else:
+            model_path = "v1/u_model.pth"
+            
+            if not os.path.exists(model_path):
+                print("❌ Model file not found at", model_path)
+                # Download the model file from GitHub
+                try:
+                    print("📥 Downloading model weights from GitHub...")
+                    import requests
+                    url = "https://github.com/malibayram/llm-from-scratch/raw/main/u_model_4000.pth"
+                    
+                    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()  # Raise an exception for bad status codes
+                    
+                    # Check if we got a proper binary file (PyTorch files start with specific bytes)
+                    if response.content[:4] != b'PK\x03\x04' and b'<html' in response.content[:100].lower():
+                        raise Exception("Downloaded HTML instead of binary file - check URL")
+                    
+                    print(f"📦 Downloaded {len(response.content)} bytes")
+                    
+                    # Create v1 directory if it doesn't exist
+                    os.makedirs("v1", exist_ok=True)
+                    
+                    # Save the model weights to the local file system
+                    with open(model_path, "wb") as f:
+                        f.write(response.content)
+                    print("✅ Model weights saved successfully!")
+                except Exception as e:
+                    print(f"❌ Failed to download model weights: {e}")
+                    print("Using random initialization.")
 
         if os.path.exists(model_path):
             try:
                 u_model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=False))
                 u_model.eval()
                 print("✅ Model weights loaded successfully!")
+                return u_model, u_tokenizer, f"✅ Model loaded from: {model_path}"
             except Exception as e:
                 print(f"⚠️ Warning: Could not load trained weights: {e}")
                 print("Using random initialization.")
+                return u_model, u_tokenizer, f"⚠️ Failed to load weights: {e}"
         else:
             print(f"⚠️ Model file not found at {model_path}. Using random initialization.")
-        
-        return u_model, u_tokenizer
+            return u_model, u_tokenizer, "⚠️ Using random initialization"
         
     except Exception as e:
         print(f"❌ Error loading model: {e}")
         raise e
 
+# Global model variables
+model, tokenizer, model_status = None, None, "Not loaded"
+
 # Initialize model and tokenizer globally
 try:
-    model, tokenizer = load_model()
+    model, tokenizer, model_status = load_model()
     print("🚀 UstaModel and tokenizer initialized successfully!")
 except Exception as e:
     print(f"❌ Failed to initialize model: {e}")
-    model, tokenizer = None, None
+    model, tokenizer, model_status = None, None, f"❌ Error: {e}"
+
+def update_model(uploaded_file):
+    """Update the model when a new file is uploaded"""
+    global model, tokenizer, model_status
+    
+    if uploaded_file is None:
+        return "❌ No file uploaded"
+    
+    try:
+        # Load the new model
+        new_model, new_tokenizer, status = load_model(uploaded_file.name)
+        
+        # Update global variables
+        model = new_model
+        tokenizer = new_tokenizer
+        model_status = status
+        
+        return status
+    except Exception as e:
+        error_msg = f"❌ Failed to load uploaded model: {e}"
+        model_status = error_msg
+        return error_msg
 
 def respond(
     message,
@@ -145,28 +175,65 @@ def respond(
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(
-            value="You are Usta, a geographical knowledge assistant trained from scratch.", 
-            label="System message",
-            info="Note: This model focuses on geographical knowledge (countries, capitals, cities)"
-        ),
-        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)",
-            info="Note: This parameter is not used by UstaModel but kept for interface compatibility"
-        ),
-    ],
-    title="🤖 Usta Model Chat",
-    description="Chat with a custom transformer language model built from scratch! This model specializes in geographical knowledge including countries, capitals, and cities."
-)
+
+# 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"],
+                    info="Upload a custom UstaModel checkpoint to use instead of the default model"
+                )
+                upload_btn = gr.Button("Load Model", variant="primary")
+                model_status_display = gr.Textbox(
+                    label="Model Status",
+                    value=model_status,
+                    interactive=False,
+                    info="Shows the current model loading status"
+                )
+        
+        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",
+                    info="Note: This model focuses on geographical knowledge"
+                )
+                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)",
+                    info="Note: This parameter is not used by UstaModel"
+                )
+    
+    # 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
+    )
+    
+    # Event handlers
+    upload_btn.click(
+        update_model,
+        inputs=[model_file],
+        outputs=[model_status_display]
+    )
 
 if __name__ == "__main__":
     demo.launch()