Upload openllm_training_with_auth.py with huggingface_hub

openllm_training_with_auth.py

@@ -16,38 +16,40 @@ import torch
 from pathlib import Path
 from huggingface_hub import HfApi, login, whoami, create_repo
 
+
 class OpenLLMTrainingManager:
     """Manages OpenLLM training with Hugging Face authentication."""
-    
+
     def __init__(self):
         """Initialize the training manager with authentication."""
         self.setup_authentication()
         self.api = HfApi()
         self.username = None
-        
+
     def setup_authentication(self):
         """Setup authentication using Space's built-in access token."""
         print("🔐 Setting up Space authentication...")
-        
+
         try:
             # Try Space's built-in authentication first (primary method)
             user_info = whoami()
-            self.username = user_info.get('name', 'unknown')
+            self.username = user_info.get("name", "unknown")
             print(f"✅ Space built-in authentication successful!")
             print(f"👤 User: {self.username}")
-            
+
         except Exception as e:
             print(f"❌ Space built-in authentication failed: {e}")
             print("🔄 Trying HF access token...")
-            
+
             # Fallback to HF access token
-            hf_token = os.environ.get('HF_TOKEN')
+            hf_token = os.environ.get("HF_TOKEN")
             if hf_token:
                 try:
                     from huggingface_hub import login
+
                     login(token=hf_token)
                     user_info = whoami()
-                    self.username = user_info.get('name', 'unknown')
+                    self.username = user_info.get("name", "unknown")
                     print(f"✅ HF access token authentication successful!")
                     print(f"👤 User: {self.username}")
                 except Exception as e2:
@@ -58,7 +60,7 @@ class OpenLLMTrainingManager:
                 print("❌ No authentication method available")
                 print("💡 Please set HF_TOKEN in Space settings or check Space permissions")
                 sys.exit(1)
-    
+
     def create_model_config(self, model_size="small", steps=8000):
         """Create model configuration file."""
         config = {
@@ -68,16 +70,16 @@ class OpenLLMTrainingManager:
             "framework": "pytorch",
             "license": "GPL-3.0",
             "author": "Louis Chua Bean Chong",
-            "description": f"OpenLLM {model_size} model trained for {steps} steps"
+            "description": f"OpenLLM {model_size} model trained for {steps} steps",
         }
-        
+
         config_path = Path("model_config.json")
-        with open(config_path, 'w') as f:
+        with open(config_path, "w") as f:
             json.dump(config, f, indent=2)
-        
+
         print(f"✅ Model config created: {config_path}")
         return config_path
-    
+
     def create_model_card(self, model_size="small", steps=8000):
         """Create model card README."""
         readme_content = f"""# OpenLLM {model_size.title()} Model
@@ -108,36 +110,31 @@ Louis Chua Bean Chong
 
 GPL-3.0
 """
-        
+
         readme_path = Path("README.md")
-        with open(readme_path, 'w') as f:
+        with open(readme_path, "w") as f:
             f.write(readme_content)
-        
+
         print(f"✅ Model card created: {readme_path}")
         return readme_path
-    
+
     def upload_model(self, model_dir, model_size="small", steps=8000):
         """Upload trained model to Hugging Face Hub."""
         print(f"📤 Uploading model to Hugging Face Hub...")
-        
+
         # Create model repository name
         repo_name = f"openllm-{model_size}-{steps}steps"
         repo_id = f"{self.username}/{repo_name}"
-        
+
         try:
             # Create repository
             print(f"🔄 Creating repository: {repo_id}")
-            create_repo(
-                repo_id=repo_id,
-                repo_type="model",
-                exist_ok=True,
-                private=False
-            )
-            
+            create_repo(repo_id=repo_id, repo_type="model", exist_ok=True, private=False)
+
             # Create model files
             config_path = self.create_model_config(model_size, steps)
             readme_path = self.create_model_card(model_size, steps)
-            
+
             # Upload files
             print(f"📁 Uploading model files...")
             self.api.upload_file(
@@ -145,17 +142,17 @@ GPL-3.0
                 path_in_repo="config.json",
                 repo_id=repo_id,
                 repo_type="model",
-                commit_message="Add model configuration"
+                commit_message="Add model configuration",
             )
-            
+
             self.api.upload_file(
                 path_or_fileobj=str(readme_path),
                 path_in_repo="README.md",
                 repo_id=repo_id,
                 repo_type="model",
-                commit_message="Add model card"
+                commit_message="Add model card",
             )
-            
+
             # Upload model files if they exist
             model_path = Path(model_dir)
             if model_path.exists():
@@ -164,17 +161,17 @@ GPL-3.0
                     folder_path=model_dir,
                     repo_id=repo_id,
                     repo_type="model",
-                    commit_message=f"Add OpenLLM {model_size} model ({steps} steps)"
+                    commit_message=f"Add OpenLLM {model_size} model ({steps} steps)",
                 )
-            
+
             print(f"✅ Model uploaded successfully!")
             print(f"🔗 Model URL: https://huggingface.co/{repo_id}")
             return repo_id
-            
+
         except Exception as e:
             print(f"❌ Model upload failed: {e}")
             return None
-    
+
     def run_training(self, model_size="small", steps=8000):
         """Run the OpenLLM training process."""
         print(f"🚀 Starting OpenLLM Training")
@@ -182,46 +179,46 @@ GPL-3.0
         print(f"📊 Model Size: {model_size}")
         print(f"🔄 Training Steps: {steps}")
         print(f"👤 User: {self.username}")
-        
+
         # Simulate training process
         print(f"\n🔄 Step 1: Initializing training...")
         print(f"   - Setting up PyTorch environment")
         print(f"   - Loading training data")
         print(f"   - Configuring model architecture")
-        
+
         print(f"\n🔄 Step 2: Training model...")
         for step in range(1, min(steps + 1, 11)):  # Show first 10 steps
             loss = 6.5 - (step * 0.1)  # Simulate decreasing loss
-            lr = 0.001 * (0.95 ** step)  # Simulate learning rate decay
+            lr = 0.001 * (0.95**step)  # Simulate learning rate decay
             print(f"   Step {step}/{steps} | Loss: {loss:.4f} | LR: {lr:.2e}")
-        
+
         if steps > 10:
             print(f"   ... (showing first 10 steps)")
             print(f"   Final step {steps} | Loss: {6.5 - (steps * 0.1):.4f}")
-        
+
         print(f"\n🔄 Step 3: Saving model...")
         model_dir = f"./openllm-trained-{model_size}"
         os.makedirs(model_dir, exist_ok=True)
-        
+
         # Create dummy model files
         model_files = [
             "best_model.pt",
             "checkpoint_step_1000.pt",
             "tokenizer/tokenizer.model",
-            "config.json"
+            "config.json",
         ]
-        
+
         for file_name in model_files:
             file_path = Path(model_dir) / file_name
             file_path.parent.mkdir(parents=True, exist_ok=True)
-            with open(file_path, 'w') as f:
+            with open(file_path, "w") as f:
                 f.write(f"# Dummy {file_name} file for demonstration")
-        
+
         print(f"✅ Model saved to: {model_dir}")
-        
+
         print(f"\n🔄 Step 4: Uploading model...")
         repo_id = self.upload_model(model_dir, model_size, steps)
-        
+
         if repo_id:
             print(f"\n🎉 Training completed successfully!")
             print(f"📊 Results:")
@@ -232,35 +229,37 @@ GPL-3.0
         else:
             print(f"\n❌ Training completed but upload failed")
             print(f"   - Model saved locally: {model_dir}")
-        
+
         return repo_id
 
+
 def main():
     """Main function to run OpenLLM training."""
     print("🚀 OpenLLM Training with Space Authentication")
     print("=" * 55)
-    
+
     # Initialize training manager
     try:
         manager = OpenLLMTrainingManager()
     except Exception as e:
         print(f"❌ Failed to initialize training manager: {e}")
         sys.exit(1)
-    
+
     # Run training
     try:
         repo_id = manager.run_training(model_size="small", steps=8000)
-        
+
         if repo_id:
             print(f"\n✅ Training and upload completed successfully!")
             print(f"🚀 Your model is ready at: https://huggingface.co/{repo_id}")
         else:
             print(f"\n⚠️ Training completed but upload failed")
             print(f"🔧 Check authentication and try again")
-            
+
     except Exception as e:
         print(f"❌ Training failed: {e}")
         sys.exit(1)
 
+
 if __name__ == "__main__":
     main()