Update hf_merge.py

hf_merge.py

@@ -3,7 +3,7 @@ import shutil
 import argparse
 import requests
 from tqdm import tqdm
-from huggingface_hub import HfApi, hf_hub_download, create_repo
+from huggingface_hub import HfApi, hf_hub_download
 from merge import merge_folder, map_tensors_to_files, copy_nontensor_files, save_tensor_map
 import logging
 
@@ -47,7 +47,6 @@ class RepositoryManager:
                 file_url = f"https://huggingface.co/{repo_name}/resolve/main/{file_path}"
                 hf_hub_download(repo_id=repo_name, filename=file_path, cache_dir=path, local_dir=path)
 
-
     def delete_repo(self, path):
         """
         Delete a repository from the local filesystem.
@@ -114,24 +113,26 @@ class ModelMerger:
         copy_nontensor_files(base_model_path, output_dir)
         save_tensor_map(self.tensor_map, output_dir)
 
-    
-    
     def upload_model(self, output_dir, commit_message):
-            """
-            Upload the merged model to HuggingFace.
-    
-            Args:
-                output_dir (str): The path to the output directory containing the merged model.
-                commit_message (str): The commit message for the upload.
-            """
-    
-                # Upload the folder to the repository
-        self.api.upload_folder(
+        """
+        Upload the merged model to HuggingFace.
+
+        Args:
+            output_dir (str): The path to the output directory containing the merged model.
+            commit_message (str): The commit message for the upload.
+        """
+        try:
+            # Create a new repository if it doesn't exist
+            if not self.api.repo_exists(repo_id=self.repo_id, token=self.token):
+                self.api.create_repo(repo_id=self.repo_id, token=self.token, private=True)
+
+            # Upload the folder to the repository
+            self.api.upload_folder(
                 repo_id=self.repo_id,
                 folder_path=output_dir,
                 commit_message=commit_message,
                 token=self.token
-                )
+            )
             url = f"https://huggingface.co/{self.repo_id}"
             logging.info(f"Model uploaded successfully to {url}")
             return url
@@ -139,6 +140,7 @@ class ModelMerger:
             logging.error(f"Error: Failed to upload to {self.repo_id}.")
             logging.error(e)
             return ""
+
 def get_max_vocab_size(repo_list):
     """
     Get the maximum vocabulary size from a list of repositories.
@@ -222,7 +224,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-
-if __name__ == "__main__":
-    main()