Update app.py

app.py

@@ -71,15 +71,16 @@ def push_existing_model(model_name):
         output += "✅ Login successful!\n\n"
         
         output += "📂 Checking for trained model...\n"
-        # 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."
+        # FIXED: Script 2 saves to ./cuad_lora_out, not ./model_output
+        model_dir = "./cuad_lora_out"
+        if not os.path.exists(model_dir):
+            return output + f"❌ {model_dir} directory not found!\nMake sure you've run training first."
         
         # 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 += f"📂 Loading trained model from {model_dir}...\n"
+            model = AutoPeftModelForQuestionAnswering.from_pretrained(model_dir)
+            tokenizer = AutoTokenizer.from_pretrained(model_dir)
             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."
@@ -95,21 +96,6 @@ def push_existing_model(model_name):
         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):
-            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:
@@ -131,7 +117,7 @@ with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
             **Instructions:**
             1. Enter your desired model name (format: `AvocadoMuffin/model-name`)
             2. Click "Start Training"
-            3. Wait ~45-75 minutes for training to complete
+            3. Wait ~25-30 minutes for training to complete
             
             **Example model names:**
             - `AvocadoMuffin/roberta-cuad-qa`
@@ -151,9 +137,9 @@ with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
                 with gr.Column():
                     gr.Markdown("""
                     **Training Details:**
-                    - Model: RoBERTa-base + LoRA
-                    - Dataset: CUAD (4000 samples → ~20K after tokenization)
-                    - Time: ~45-75 minutes
+                    - Model: RoBERTa-base-squad2 + LoRA
+                    - Dataset: CUAD (balanced, ~18K samples)
+                    - Time: ~25-30 minutes on L4/T4
                     - GPU: T4 (free tier)
                     - Will auto-push to your HF profile
                     """)
@@ -191,7 +177,7 @@ with gr.Blocks(title="RoBERTa CUAD Trainer") as demo:
                 with gr.Column():
                     gr.Markdown("""
                     **Requirements:**
-                    - Model must be trained and saved in ./model_output/
+                    - Model must be trained and saved in ./cuad_lora_out/
                     - roberta_token must be set in Space secrets
                     - Takes ~30 seconds
                     """)