Update app.py

app.py

@@ -74,38 +74,46 @@ def push_existing_model(model_name):
         # Check if model exists
         if not os.path.exists("./model_output"):
             return output + "❌ ./model_output directory not found!\nMake sure you've run training first."
-            
-        # Load your already-trained model
-        output += "📂 Loading trained model from ./model_output...\n"
-        model = AutoPeftModelForQuestionAnswering.from_pretrained("./model_output")
-        tokenizer = AutoTokenizer.from_pretrained("./model_output")
-        output += "✅ Model loaded successfully!\n\n"
+        
+        # FIXED: Better error handling for model loading
+        try:
+            output += "📂 Loading trained model from ./model_output...\n"
+            model = AutoPeftModelForQuestionAnswering.from_pretrained("./model_output")
+            tokenizer = AutoTokenizer.from_pretrained("./model_output")
+            output += "✅ Model loaded successfully!\n\n"
+        except Exception as e:
+            return output + f"❌ Failed to load model: {str(e)}\nMake sure the model was trained successfully."
         
         # Push to Hub
         model_name = model_name.strip()
         output += f"⬆️ Pushing model to Hub: {model_name}\n"
         
-        model.push_to_hub(model_name, private=False)
-        tokenizer.push_to_hub(model_name, private=False)
-        
-        output += f"🎉 SUCCESS! Model pushed to: https://huggingface.co/{model_name}\n"
+        try:
+            model.push_to_hub(model_name, private=False)
+            tokenizer.push_to_hub(model_name, private=False)
+            output += f"🎉 SUCCESS! Model pushed to: https://huggingface.co/{model_name}\n"
+        except Exception as e:
+            return output + f"❌ Failed to push model: {str(e)}\nCheck if model name already exists or token has write permissions."
         
         # Also push training info if it exists
         training_info_path = "./model_output/training_info.json"
         if os.path.exists(training_info_path):
-            from huggingface_hub import upload_file
-            upload_file(
-                path_or_fileobj=training_info_path,
-                path_in_repo="training_info.json",
-                repo_id=model_name,
-                repo_type="model"
-            )
-            output += "📊 Training info also uploaded!\n"
+            try:
+                from huggingface_hub import upload_file
+                upload_file(
+                    path_or_fileobj=training_info_path,
+                    path_in_repo="training_info.json",
+                    repo_id=model_name,
+                    repo_type="model"
+                )
+                output += "📊 Training info also uploaded!\n"
+            except Exception as e:
+                output += f"⚠️ Training info upload failed: {str(e)}\n"
         
         return output
         
     except Exception as e:
-        return f"❌ Error: {str(e)}\n\nCommon issues:\n- Invalid token\n- Model name already exists\n- Network issues"
+        return f"❌ Error: {str(e)}\n\nCommon issues:\n- Invalid token\n- Model name already exists\n- Network issues\n- Token lacks write permissions"
 
 # Gradio Interface
 with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
@@ -113,23 +121,30 @@ with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
     # 🤖 RoBERTa CUAD Question Answering Trainer
     
     Train a RoBERTa model with LoRA on CUAD dataset OR push an existing trained model to Hub.
+    
+    **For AvocadoMuffin:** Your models will be saved as `AvocadoMuffin/your-model-name`
     """)
     
     with gr.Tabs():
         with gr.TabItem("🚀 Train New Model"):
             gr.Markdown("""
             **Instructions:**
-            1. Enter your desired model name (format: `your-username/model-name`)
+            1. Enter your desired model name (format: `AvocadoMuffin/model-name`)
             2. Click "Start Training"
             3. Wait ~45-75 minutes for training to complete
+            
+            **Example model names:**
+            - `AvocadoMuffin/roberta-cuad-qa`
+            - `AvocadoMuffin/roberta-legal-qa-v1`
             """)
             
             with gr.Row():
                 with gr.Column():
                     model_name_input = gr.Textbox(
                         label="Model Name",
-                        placeholder="your-username/roberta-cuad-qa",
-                        info="This will be your model's name on Hugging Face Hub"
+                        placeholder="AvocadoMuffin/roberta-cuad-qa",
+                        info="This will be your model's name on Hugging Face Hub",
+                        value="AvocadoMuffin/"
                     )
                     train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
                     
@@ -137,9 +152,10 @@ with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
                     gr.Markdown("""
                     **Training Details:**
                     - Model: RoBERTa-base + LoRA
-                    - Dataset: CUAD (4000 samples → 20K after tokenization)
+                    - Dataset: CUAD (4000 samples → ~20K after tokenization)
                     - Time: ~45-75 minutes
-                    - GPU: T4 (free)
+                    - GPU: T4 (free tier)
+                    - Will auto-push to your HF profile
                     """)
             
             train_output = gr.Textbox(
@@ -166,8 +182,9 @@ with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
                 with gr.Column():
                     push_model_name = gr.Textbox(
                         label="Model Name",
-                        placeholder="your-username/roberta-cuad-qa",
-                        info="Name for your model on Hugging Face Hub"
+                        placeholder="AvocadoMuffin/roberta-cuad-qa",
+                        info="Name for your model on Hugging Face Hub",
+                        value="AvocadoMuffin/"
                     )
                     push_btn = gr.Button("⬆️ Push to Hub", variant="secondary", size="lg")
                     
@@ -194,9 +211,10 @@ with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
     
     gr.Markdown("""
     ---
-    **Setup Required:**
-    - Set `roberta_token` in Space Settings → Repository secrets
-    - Get your token from: https://huggingface.co/settings/tokens (with Write permissions)
+    **Setup Required for AvocadoMuffin:**
+    - ✅ Set `roberta_token` in Space Settings → Repository secrets  
+    - ✅ Get your token from: https://huggingface.co/settings/tokens (with Write permissions)
+    - ✅ Your trained models will appear at: `https://huggingface.co/AvocadoMuffin/model-name`
     """)
 
 if __name__ == "__main__":