""" 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)