Update app.py

app.py

@@ -1,235 +1,151 @@
-import os
-
-os.system('pip install huggingface_hub')
-
-
-import shutil
-import time
+import logging
+from pathlib import Path
+from typing import List, Optional
 
 import gradio as gr
 from huggingface_hub import HfApi
-
-CACHE_DIR = "hf2ms_cache"
-LOCAL_DIR = "hf2ms_local"
-USERNAME = "dailingx"
-EMAIL = "[email protected]"
-
-
-def clone_from_ms(
-        ms_token: str,
-        ms_repo_id: str,
-        clone_dir: str
-):
-    if os.path.exists(clone_dir):
-        ori_dir = os.getcwd()
-        os.chdir(clone_dir)
-        os.system("GIT_LFS_SKIP_SMUDGE=1 git pull")
-        os.chdir(ori_dir)
-        message = f"the repo already exists, so just pull successfully!"
-    command = f"GIT_LFS_SKIP_SMUDGE=1 git clone https://oauth2:{ms_token}@www.modelscope.cn/datasets/{ms_repo_id}.git"
-    os.system(command)
-    configuration_path = os.path.join(clone_dir, "configuration.json")
-    if not os.path.exists(configuration_path):
-        open(configuration_path, "w").close()
-    # message
-    message = f"clone from https://oauth2:{ms_token}@www.modelscope.cn/datasets/{ms_repo_id}.git successfully!"
-    return message
-
-
-def hf_list_repo_files(
-        hf_token: str,
-        hf_repo_id: str,
-        repo_type: str = "dataset"
-):
-    hf_api = HfApi(token=hf_token)
-    files = hf_api.list_repo_files(repo_id=hf_repo_id, repo_type=repo_type)
-    return files
-
-
-def pull_from_hf(
-        hf_token: str,
-        hf_repo_id: str,
-        filename: str,
-):
-    if not hf_repo_id:
-        raise gr.Error("Please enter the repo_id of huggingface")
-    if not filename:
-        raise gr.Error("Please enter the filename")
-
-    if "," in filename:
-        filename_list = filename.split(",")
-        for _filename in filename_list:
-            save_path = os.path.join(LOCAL_DIR, _filename)
-            if os.path.exists(save_path):
-                message = "the file already exists!"
-                return message
-
-            # download
-            hf_api = HfApi(token=hf_token)
-            hf_api.hf_hub_download(
-                repo_id=hf_repo_id,
+from modelscope.hub.api import HubApi
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler()]
+)
+logger = logging.getLogger(__name__)
+
+class HFToMSConverter:
+    def __init__(self, config: dict):
+        self.config = config
+        self.cache_dir = config.get('cache_dir', "hf2ms_cache")
+        self.local_dir = config.get('local_dir', "hf2ms_local")
+        self.hf_api = HfApi(token=config['hf_token'])
+        self.ms_api = HubApi()
+        self.ms_api.login(config['ms_token'])
+        
+        for dir_path in [self.local_dir, self.cache_dir]:
+            Path(dir_path).mkdir(exist_ok=True)
+
+    def get_hf_files(self, repo_id: str, repo_type: str = "dataset") -> List[str]:
+        """获取HuggingFace仓库文件列表"""
+        return self.hf_api.list_repo_files(repo_id=repo_id, repo_type=repo_type)
+
+    def download_file(self, repo_id: str, filename: str) -> Optional[str]:
+        """从HuggingFace下载文件"""
+        save_path = Path(self.local_dir) / filename
+        if save_path.exists():
+            logger.warning(f"文件已存在: {filename}")
+            return None
+
+        try:
+            self.hf_api.hf_hub_download(
+                repo_id=repo_id,
                 repo_type="dataset",
-                filename=_filename,
-                cache_dir=CACHE_DIR,
-                local_dir=LOCAL_DIR,
+                filename=filename,
+                cache_dir=self.cache_dir,
+                local_dir=self.local_dir,
                 local_dir_use_symlinks=False
             )
-    else:
-        save_path = os.path.join(LOCAL_DIR, filename)
-        if os.path.exists(save_path):
-            message = "the file already exists!"
-            return message
-
-        # download
-        hf_api = HfApi(token=hf_token)
-        hf_api.hf_hub_download(
-            repo_id=hf_repo_id,
-            repo_type="dataset",
-            filename=filename,
-            cache_dir=CACHE_DIR,
-            local_dir=LOCAL_DIR,
-            local_dir_use_symlinks=False
+            logger.info(f"成功下载文件: {filename}")
+            return str(save_path)
+        except Exception as e:
+            logger.error(f"下载失败 {filename}: {e}")
+            return None
+
+    def handle_file_operation(self, operation_type: str, *args) -> bool:
+        """统一处理文件操作的异常"""
+        try:
+            if operation_type == "move":
+                src, dst = args
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                src.rename(dst)
+                logger.info(f"移动文件成功: {src.name}")
+            elif operation_type == "push":
+                ms_repo_id, clone_dir, filename = args
+                self.ms_api.upload_folder(
+                    repo_id=f"{self.config['username']}/{ms_repo_id}",
+                    folder_path=clone_dir,
+                    commit_message='upload dataset folder',
+                    repo_type='dataset'
+                )
+                logger.info(f"推送成功: {filename}")
+            return True
+        except Exception as e:
+            logger.error(f"{operation_type}操作失败: {e}")
+            return False
+
+    def move_file(self, filename: str, clone_dir: str) -> bool:
+        """移动文件到目标目录"""
+        return self.handle_file_operation(
+            "move",
+            Path(self.local_dir) / filename,
+            Path(clone_dir) / filename
         )
 
-    # message
-    message = f"Pull from https://huggingface.co/datasets/{hf_repo_id} successfully!"
-    print(message)
-    return message
-
-
-def move_file_from_local_to_clone_dir(
-        filename: str,
-        clone_dir: str
-):
-    # move to the clone dir
-    if "," in filename:
-        filename_list = filename.split(",")
-        for _filename in filename_list:
-            if "/" in _filename:
-                dirname = os.path.dirname(_filename)
-                src_dir = os.path.join(clone_dir, dirname)
-                if not os.path.exists(src_dir):
-                    os.makedirs(src_dir)
-            src = os.path.join(LOCAL_DIR, _filename)
-            dst = os.path.join(clone_dir, _filename)
-            shutil.move(src=src, dst=dst)
-    else:
-        if "/" in filename:
-            dirname = os.path.dirname(filename)
-            src_dir = os.path.join(clone_dir, dirname)
-            if not os.path.exists(src_dir):
-                os.makedirs(src_dir)
-        src = os.path.join(LOCAL_DIR, filename)
-        dst = os.path.join(clone_dir, filename)
-        shutil.move(src=src, dst=dst)
-
-    # message
-    message = f"move the file from {src} to {dst} successfully!"
-    print(message)
-    return message
-
-
-def push_to_ms(
-        username: str,
-        email: str,
-        ms_repo_id: str,
-        clone_dir: str,
-        filename: str
-):
-    # push to ms
-    ori_dir = os.getcwd()
-    os.chdir(clone_dir)
-    os.system("apt-get install git-lfs")
-    os.system("git lfs install")
-    os.system(f"git config --global user.email {email}")
-    os.system(f"git config --global user.name {username}")
-    print(filename)
-    
-    if "," in filename:
-        filename_list = filename.split(",")
-        num = len(filename_list)
-        for _filename in filename_list:
-            file_path = _filename.replace(os.path.basename(_filename), '')
-            print(file_path)
-            os.system("cd {file_path}")
-            os.system(f"git lfs track *")
-            print(file_path)
-            current_working_directory = os.getcwd()
-            print(current_working_directory)
-            os.system('cd {ori_dir}')
-
-        os.system(f"git lfs track *")
-        os.system("git add .")
-        os.system(f"git commit -m 'upload {num} files'")
-        os.system(f"git push")
-        os.chdir(ori_dir)
-    else:
-        os.system(f"git lfs track '{filename}'")
-        os.system("git add .")
-        os.system(f"git commit -m 'upload {filename}'")
-        os.system(f"git push")
-        os.chdir(ori_dir)
-    # remove clone dir
-    if os.path.exists(clone_dir):
-        shutil.rmtree(clone_dir)
-    message = f'Pushed to https://www.modelscope.cn/datasets/{ms_repo_id} successfully!'
-    print(message)
-    return message
-
-
-def handle(
-        hf_token: str,
-        ms_token: str,
-        repo_type: str,
-        hf_repo: str,
-        ms_repo: str,
-):
-    clone_dir = ms_repo.split("/")[-1]
-    hf_file_list = hf_list_repo_files(hf_token, hf_repo, repo_type)
-    print(f"all file in hf: {hf_file_list}")
-
-    for filename in hf_file_list:
-        print(f"begin to process file: {filename}")
-        clone_from_ms(ms_token, ms_repo, clone_dir)
-        time.sleep(1)
-        pull_from_hf(hf_token, hf_repo, filename)
-        time.sleep(1)
-        move_file_from_local_to_clone_dir(filename, clone_dir)
-        time.sleep(1)
-        push_to_ms(USERNAME, EMAIL, ms_repo, clone_dir, filename)
-        print(f"process file finish: {filename}")
-        time.sleep(10)
-
-
-with gr.Blocks() as demo:
-    gr.Markdown(
-        '''
-        This space uploads model from Huggingface to ModelScope.
-        **Please make sure that you're the owner of the repo or have permission from the owner to do so!**
-        # How to use this Space?
-        - Duplicate this Space and providing MS token (optional) and your read/write HF token (mandatory)
-        - Create your target model repo on HF. This step needs to be done manually. The Space doesn't do create an empty repo for you.
-        - In your own private Space, fill in information below.
-        - Click submit then watch for output in container log for progress.
-        - Create README.md file (since the metadata is not compatible with HF)
-        '''
-    )
-    hf_token = gr.Textbox(label="HuggingFace Token")
-    ms_token = gr.Textbox(label="ModelScope Git Token")
-    repo_type = gr.Textbox(label="Repo Type", value="dataset")
-    hf_repo = gr.Textbox(label="HuggingFace Repo")
-    ms_repo = gr.Textbox(label="ModelScope Repo")
-
-    with gr.Row():
-        button = gr.Button("Submit", variant="primary")
-        clear = gr.Button("Clear")
+    def push_to_ms(self, ms_repo_id: str, clone_dir: str, filename: str) -> bool:
+        """推送到ModelScope"""
+        return self.handle_file_operation("push", ms_repo_id, clone_dir, filename)
+
+    def process_file(self, hf_repo: str, ms_repo: str, filename: str) -> bool:
+        """处理单个文件的完整流程"""
+        clone_dir = ms_repo.split("/")[-1]
+        
+        try:
+            if not all([
+                self.download_file(hf_repo, filename),
+                self.move_file(filename, clone_dir),
+                self.push_to_ms(ms_repo, clone_dir, filename)
+            ]):
+                return False
+            return True
+        except Exception as e:
+            logger.error(f"处理文件失败 {filename}: {e}")
+            return False
+
+def create_ui() -> gr.Blocks:
+    """创建Gradio界面"""
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            """
+            # HuggingFace to ModelScope 数据迁移工具
+            请确保您拥有相应仓库的权限。
+            """
+        )
+        
+        with gr.Row():
+            hf_token = gr.Textbox(label="HuggingFace Token")
+            ms_token = gr.Textbox(label="ModelScope Token")
+            
+        with gr.Row():
+            repo_type = gr.Textbox(label="仓库类型", value="dataset")
+            hf_repo = gr.Textbox(label="HuggingFace仓库")
+            ms_repo = gr.Textbox(label="ModelScope仓库")
+            
+        with gr.Row():
+            submit = gr.Button("开始迁移", variant="primary")
+            clear = gr.Button("清除")
+
+        def handle_submit(hf_token, ms_token, repo_type, hf_repo, ms_repo):
+            config = {
+                'hf_token': hf_token,
+                'ms_token': ms_token,
+                'username': "thomas",
+                'email': "yx20001210@163.com",
+            }
+            
+            converter = HFToMSConverter(config)
+            files = converter.get_hf_files(hf_repo, repo_type)
+            
+            for filename in files:
+                converter.process_file(hf_repo, ms_repo, filename)
+
+        submit.click(
+            handle_submit,
+            inputs=[hf_token, ms_token, repo_type, hf_repo, ms_repo],
+        )
 
-    button.click(
-        handle,
-        [hf_token, ms_token, repo_type, hf_repo, ms_repo],
-        outputs=None
-    )
+    return demo
 
 if __name__ == "__main__":
+    demo = create_ui()
     demo.queue(max_size=1)
     demo.launch(share=False, max_threads=1)