Add HuggingFace Hub service

app/services/hub_service.py

@@ -0,0 +1,267 @@
+"""
+HuggingFace Hub Service - Handle model pushing and repository management
+"""
+
+import logging
+from typing import Dict, Any, Optional, List
+from huggingface_hub import (
+    HfApi,
+    create_repo,
+    upload_folder,
+    upload_file,
+    delete_repo,
+    repo_info,
+    list_repo_files,
+)
+import os
+import tempfile
+import json
+
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceHubService:
+    """Service for interacting with HuggingFace Hub."""
+    
+    def __init__(self, token: Optional[str] = None):
+        self.token = token
+        self.api = HfApi(token=token)
+    
+    def create_model_repo(
+        self,
+        repo_name: str,
+        private: bool = False,
+        exist_ok: bool = True,
+    ) -> Dict[str, Any]:
+        """Create a new model repository."""
+        try:
+            url = create_repo(
+                repo_id=repo_name,
+                token=self.token,
+                private=private,
+                exist_ok=exist_ok,
+            )
+            return {
+                "success": True,
+                "repo_id": repo_name,
+                "url": url,
+                "private": private,
+            }
+        except Exception as e:
+            logger.error(f"Error creating repo {repo_name}: {e}")
+            return {
+                "success": False,
+                "error": str(e),
+            }
+    
+    def push_model(
+        self,
+        model_path: str,
+        repo_id: str,
+        commit_message: str = "Push model via Universal Trainer",
+        commit_description: Optional[str] = None,
+        private: bool = False,
+    ) -> Dict[str, Any]:
+        """Push a model to the HuggingFace Hub."""
+        try:
+            # Create repo if it doesn't exist
+            self.create_model_repo(repo_id, private=private)
+            
+            # Upload the model folder
+            result = upload_folder(
+                folder_path=model_path,
+                repo_id=repo_id,
+                token=self.token,
+                commit_message=commit_message,
+                commit_description=commit_description,
+            )
+            
+            return {
+                "success": True,
+                "repo_id": repo_id,
+                "url": f"https://huggingface.co/{repo_id}",
+                "commit_url": result.commit_url if hasattr(result, 'commit_url') else None,
+            }
+            
+        except Exception as e:
+            logger.error(f"Error pushing model to {repo_id}: {e}")
+            return {
+                "success": False,
+                "error": str(e),
+            }
+    
+    def push_training_artifacts(
+        self,
+        output_dir: str,
+        repo_id: str,
+        model_name: str,
+        task_type: str,
+        dataset_name: str,
+        training_config: Dict[str, Any],
+        metrics: Dict[str, float],
+    ) -> Dict[str, Any]:
+        """Push complete training artifacts to Hub."""
+        try:
+            # Create repo
+            create_result = self.create_model_repo(repo_id)
+            
+            # Create/update model card
+            model_card = self._create_model_card(
+                model_name=model_name,
+                task_type=task_type,
+                dataset_name=dataset_name,
+                training_config=training_config,
+                metrics=metrics,
+            )
+            
+            # Write model card
+            readme_path = os.path.join(output_dir, "README.md")
+            with open(readme_path, "w") as f:
+                f.write(model_card)
+            
+            # Upload everything
+            result = upload_folder(
+                folder_path=output_dir,
+                repo_id=repo_id,
+                token=self.token,
+                commit_message=f"Upload fine-tuned {model_name}",
+            )
+            
+            return {
+                "success": True,
+                "repo_id": repo_id,
+                "url": f"https://huggingface.co/{repo_id}",
+            }
+            
+        except Exception as e:
+            logger.error(f"Error pushing artifacts: {e}")
+            return {
+                "success": False,
+                "error": str(e),
+            }
+    
+    def delete_model(self, repo_id: str) -> Dict[str, Any]:
+        """Delete a model repository."""
+        try:
+            delete_repo(repo_id=repo_id, token=self.token)
+            return {
+                "success": True,
+                "message": f"Deleted {repo_id}",
+            }
+        except Exception as e:
+            logger.error(f"Error deleting repo {repo_id}: {e}")
+            return {
+                "success": False,
+                "error": str(e),
+            }
+    
+    def get_repo_info(self, repo_id: str) -> Dict[str, Any]:
+        """Get repository information."""
+        try:
+            info = repo_info(repo_id=repo_id, token=self.token)
+            return {
+                "success": True,
+                "repo_id": info.id,
+                "author": info.author,
+                "private": info.private,
+                "downloads": getattr(info, "downloads", 0),
+                "likes": getattr(info, "likes", 0),
+                "tags": info.tags or [],
+                "sha": info.sha,
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "error": str(e),
+            }
+    
+    def list_model_files(self, repo_id: str) -> Dict[str, Any]:
+        """List files in a model repository."""
+        try:
+            files = list_repo_files(repo_id=repo_id, token=self.token)
+            return {
+                "success": True,
+                "files": files,
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "error": str(e),
+            }
+    
+    def _create_model_card(
+        self,
+        model_name: str,
+        task_type: str,
+        dataset_name: str,
+        training_config: Dict[str, Any],
+        metrics: Dict[str, float],
+    ) -> str:
+        """Create a model card README."""
+        
+        metrics_str = "\n".join([f"- **{k}:** {v:.4f}" if isinstance(v, float) else f"- **{k}:** {v}" for k, v in metrics.items()])
+        
+        return f"""---
+license: apache-2.0
+base_model: {model_name}
+tags:
+- {task_type}
+- fine-tuned
+- universal-trainer
+---
+
+# {model_name} Fine-tuned
+
+This model is a fine-tuned version of [{model_name}](https://huggingface.co/{model_name}) on the [{dataset_name}](https://huggingface.co/datasets/{dataset_name}) dataset.
+
+## Model Details
+
+- **Base Model:** {model_name}
+- **Task Type:** {task_type}
+- **Dataset:** {dataset_name}
+
+## Training Configuration
+
+| Parameter | Value |
+|-----------|-------|
+| Learning Rate | {training_config.get('learning_rate', 'N/A')} |
+| Batch Size | {training_config.get('batch_size', 'N/A')} |
+| Epochs | {training_config.get('epochs', 'N/A')} |
+| Max Length | {training_config.get('max_length', 'N/A')} |
+| PEFT | {training_config.get('use_peft', False)} |
+
+## Training Metrics
+
+{metrics_str if metrics else 'No metrics available'}
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("{model_name}")
+tokenizer = AutoTokenizer.from_pretrained("{model_name}")
+```
+
+## Training Procedure
+
+This model was trained using the [Universal Model Trainer](https://huggingface.co/spaces/vectorplasticity/universal-model-trainer).
+
+### Framework Versions
+
+- Transformers: {__import__('transformers').__version__}
+- PyTorch: {__import__('torch').__version__}
+
+## Limitations
+
+This model inherits the limitations of its base model. Please refer to the original model card for details.
+
+## License
+
+Please refer to the original model's license.
+"""
+
+
+def get_hub_client(token: Optional[str] = None) -> HuggingFaceHubService:
+    """Get a HuggingFace Hub client."""
+    return HuggingFaceHubService(token=token)