Shiyu Zhao commited on
Commit
f268b5a
·
1 Parent(s): b65f2ef

Update space

Browse files
utils/__init__.py ADDED
File without changes
utils/hub_storage.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+ from huggingface_hub import HfApi
3
+ from .token_handler import TokenHandler
4
+
5
+ class HubStorage:
6
+ def __init__(self, repo_id):
7
+ self.token_handler = TokenHandler()
8
+ self.token = self.token_handler.get_verified_token()
9
+ self.api = HfApi()
10
+ self.repo_id = repo_id
11
+
12
+ def save_to_hub(self, file_content, path_in_repo, commit_message):
13
+ """Save file to HuggingFace Hub"""
14
+ return self.api.upload_file(
15
+ path_or_fileobj=file_content,
16
+ path_in_repo=path_in_repo,
17
+ repo_id=self.repo_id,
18
+ repo_type="space",
19
+ commit_message=commit_message,
20
+ token=self.token
21
+ )
22
+
23
+ def load_from_hub(self, path_in_repo):
24
+ """Load file from HuggingFace Hub"""
25
+ return self.api.hf_hub_download(
26
+ repo_id=self.repo_id,
27
+ filename=path_in_repo,
28
+ repo_type="space",
29
+ token=self.token
30
+ )
utils/token_handler.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from dotenv import load_dotenv
3
+ from huggingface_hub import HfApi
4
+ from pathlib import Path
5
+
6
+ class TokenHandler:
7
+ def __init__(self):
8
+ # Load environment variables from .env file if it exists
9
+ self.load_environment()
10
+ self.token = self._get_token()
11
+ self.api = HfApi()
12
+
13
+ def load_environment(self):
14
+ """Load environment variables from .env file"""
15
+ env_path = Path('.env')
16
+ if env_path.exists():
17
+ load_dotenv(env_path)
18
+
19
+ def _get_token(self) -> str:
20
+ """Get HuggingFace token from environment variables"""
21
+ token = os.getenv("HF_TOKEN")
22
+ if not token:
23
+ raise EnvironmentError(
24
+ "HF_TOKEN not found in environment variables. "
25
+ "Please set it up using one of these methods:\n"
26
+ "1. Create a .env file with HF_TOKEN=your_token\n"
27
+ "2. Set environment variable HF_TOKEN=your_token\n"
28
+ "3. Add HF_TOKEN to your HuggingFace Space secrets"
29
+ )
30
+ return token
31
+
32
+ def verify_token(self) -> bool:
33
+ """Verify if the token is valid by making a test API call"""
34
+ try:
35
+ # Try to get user information using the token
36
+ self.api.whoami(token=self.token)
37
+ return True
38
+ except Exception as e:
39
+ print(f"Token verification failed: {e}")
40
+ return False
41
+
42
+ def get_verified_token(self) -> str:
43
+ """Get token and verify it's working"""
44
+ if not self.verify_token():
45
+ raise ValueError(
46
+ "Invalid or expired HuggingFace token. "
47
+ "Please check your token at https://huggingface.co/settings/tokens"
48
+ )
49
+ return self.token
50
+
51
+ # Usage example
52
+ def initialize_hf_token():
53
+ """Initialize and verify HuggingFace token"""
54
+ try:
55
+ handler = TokenHandler()
56
+ token = handler.get_verified_token()
57
+ print("✓ HuggingFace token successfully verified")
58
+ return token
59
+ except Exception as e:
60
+ print(f"✗ Error initializing HuggingFace token: {e}")
61
+ return None
62
+
63
+ # Example of how to use in your main code
64
+ if __name__ == "__main__":
65
+ # Create .env file if it doesn't exist
66
+ if not Path('.env').exists():
67
+ print("Creating .env file template...")
68
+ with open('.env', 'w') as f:
69
+ f.write("HF_TOKEN=your_token_here\n")
70
+ print("Please edit .env file and add your HuggingFace token")
71
+
72
+ # Initialize token
73
+ token = initialize_hf_token()
74
+ if token:
75
+ print("Ready to use HuggingFace API")