space update

app.py

@@ -71,7 +71,23 @@ def load_model(custom_model_path=None):
 
         if os.path.exists(model_path):
             try:
-                u_model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=False))
+                state_dict = torch.load(model_path, map_location="cpu", weights_only=False)
+                
+                # Handle potential key mapping issues
+                if "embedding.weight" in state_dict and "embedding.embedding.weight" not in state_dict:
+                    # Map old key names to new key names
+                    new_state_dict = {}
+                    for key, value in state_dict.items():
+                        if key == "embedding.weight":
+                            new_state_dict["embedding.embedding.weight"] = value
+                        elif key == "pos_embedding.weight":
+                            # Skip positional embedding if not expected
+                            continue
+                        else:
+                            new_state_dict[key] = value
+                    state_dict = new_state_dict
+                
+                u_model.load_state_dict(state_dict)
                 u_model.eval()
                 print("✅ Model weights loaded successfully!")
                 return u_model, u_tokenizer, f"✅ Model loaded from: {model_path}"
@@ -188,15 +204,13 @@ with gr.Blocks(title="🤖 Usta Model Chat", theme=gr.themes.Soft()) as demo:
                 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"
+                    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,
-                    info="Shows the current model loading status"
+                    interactive=False
                 )
         
         with gr.Column(scale=1):
@@ -205,8 +219,7 @@ with gr.Blocks(title="🤖 Usta Model Chat", theme=gr.themes.Soft()) as demo:
                 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"
+                    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")
@@ -215,8 +228,7 @@ with gr.Blocks(title="🤖 Usta Model Chat", theme=gr.themes.Soft()) as demo:
                     maximum=1.0,
                     value=0.95,
                     step=0.05,
-                    label="Top-p (nucleus sampling)",
-                    info="Note: This parameter is not used by UstaModel"
+                    label="Top-p (nucleus sampling)"
                 )
     
     # Chat interface