utils/hub_storage.py

from pathlib import Path
from huggingface_hub import HfApi
from .token_handler import TokenHandler

class HubStorage:
    def __init__(self, repo_id):
        self.token_handler = TokenHandler()
        self.token = self.token_handler.get_verified_token()
        self.api = HfApi()
        self.repo_id = repo_id

    def save_to_hub(self, file_content, path_in_repo, commit_message):
        """Save file to HuggingFace Hub"""
        return self.api.upload_file(
            path_or_fileobj=file_content,
            path_in_repo=path_in_repo,
            repo_id=self.repo_id,
            repo_type="space",
            commit_message=commit_message,
            token=self.token
        )

    def load_from_hub(self, path_in_repo):
        """Load file from HuggingFace Hub"""
        return self.api.hf_hub_download(
            repo_id=self.repo_id,
            filename=path_in_repo,
            repo_type="space",
            token=self.token
        )