Update app.py

app.py

@@ -469,11 +469,25 @@ class ArchitechAgent:
         
         try:
             login(token=hf_token)
-            model.push_to_hub(model_name, token=hf_token)
-            tokenizer.push_to_hub(model_name, token=hf_token)
-            hub_url = f"https://huggingface.co/{username}/{model_name}"
             
-            return f"""🎉 **TRAINING COMPLETE!**
+            # Try uploading with retries
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    progress(0.9 + (attempt * 0.03), f"📤 Upload attempt {attempt + 1}/{max_retries}...")
+                    
+                    # Push model with timeout
+                    model.push_to_hub(
+                        model_name, 
+                        token=hf_token,
+                        max_shard_size="500MB",
+                        safe_serialization=True
+                    )
+                    tokenizer.push_to_hub(model_name, token=hf_token)
+                    
+                    hub_url = f"https://huggingface.co/{username}/{model_name}"
+                    
+                    return f"""🎉 **TRAINING COMPLETE!**
 
 ✅ Training successful
 💾 Model saved locally
@@ -486,9 +500,47 @@ class ArchitechAgent:
 - Learning rate: {learning_rate}
 
 **Test it in the 'Test Model' tab!**"""
+                
+                except Exception as upload_error:
+                    if attempt < max_retries - 1:
+                        logger.warning(f"Upload attempt {attempt + 1} failed: {upload_error}")
+                        time.sleep(5)  # Wait before retry
+                        continue
+                    else:
+                        raise upload_error
         
         except Exception as e:
-            return f"✅ Training done but upload failed: {str(e)}\nModel saved at: {output_dir}"# ==================== GRADIO INTERFACE ====================
+            logger.error(f"Upload failed after retries: {e}")
+            
+            # Provide manual upload instructions
+            return f"""✅ **TRAINING COMPLETE!** (Upload timed out)
+
+💾 Model saved locally at: `{output_dir}`
+
+**Manual Upload Instructions:**
+1. Download your Space's files (or access via SSH if enabled)
+2. Run this command locally:
+```bash
+huggingface-cli upload {username}/{model_name} {output_dir}
+```
+
+Or use the Python API:
+```python
+from huggingface_hub import HfApi
+api = HfApi()
+api.upload_folder(
+    folder_path="{output_dir}",
+    repo_id="{username}/{model_name}",
+    token="YOUR_TOKEN"
+)
+```
+
+**Stats:**
+- Examples: {len(texts)}
+- Epochs: {num_epochs}
+- Model saved successfully!
+
+**You can still test it locally or manually upload!**"""# ==================== GRADIO INTERFACE ====================
 
 def create_gradio_interface():
     agent = ArchitechAgent()