app/services/hub_service.py

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