Update annotation.py

annotation.py

@@ -1,23 +1,26 @@
 import re
 import os
 import json
-from huggingface_hub import HfApi, HfFolder, Repository
+from huggingface_hub import HfApi, hf_hub_download
 import config
 
-# 初始化 Hub token
+# 从 Space Secret 中读取 Token
 HF_TOKEN = os.getenv("HF_TOKEN")
-REPO_ID = config.SAVE_REPO_ID
+REPO_ID = config.SAVE_REPO_ID  # 在 config.py 里配置你的 dataset repo，比如 "sunnyzjx/annotation_results"
 
-def get_user_annotation_file(username):
-    """获取用户特定的标注文件路径"""
-    os.makedirs("annotations", exist_ok=True)
+api = HfApi()
+
+
+def get_user_annotation_filename(username: str) -> str:
+    """生成用户标注文件名"""
     safe_username = re.sub(r'[\\/*?:"<>|]', "_", username)
-    return f"annotations/annotation_results_{safe_username}.json"
+    return f"annotation_results_{safe_username}.json"
+
 
 def save_annotations(username_state, annotation_results_state, tasks):
-    """保存标注结果到本地并推送到 Hub"""
+    """直接推送标注结果到 Hugging Face Hub"""
     try:
-        annotation_file = get_user_annotation_file(username_state)
+        # 组织数据
         save_data = {
             "total_tasks": len(tasks),
             "completed_tasks": len(annotation_results_state),
@@ -33,35 +36,42 @@ def save_annotations(username_state, annotation_results_state, tasks):
                 "audioB_id": f"audioB_{task_id}",
                 "username": username_state
             })
-        # 保存到本地
-        with open(annotation_file, "w", encoding="utf-8") as f:
-            json.dump(save_data, f, ensure_ascii=False, indent=2)
 
-        # 推送到 Hugging Face Hub
-        api = HfApi()
+        # 临时保存到内存/字符串
+        save_str = json.dumps(save_data, ensure_ascii=False, indent=2)
+        filename = get_user_annotation_filename(username_state)
+
+        # 上传到 Hub
         api.upload_file(
-            path_or_fileobj=annotation_file,
-            path_in_repo=os.path.basename(annotation_file),
+            path_or_fileobj=save_str.encode("utf-8"),
+            path_in_repo=filename,
             repo_id=REPO_ID,
             repo_type="dataset",
             token=HF_TOKEN
         )
 
-        return f"✅ 标注结果已保存并上传到 {REPO_ID}\n完成进度: {len(annotation_results_state)}/{len(tasks)}"
+        return f"✅ 标注结果已上传到 {REPO_ID}/{filename}\n完成进度: {len(annotation_results_state)}/{len(tasks)}"
 
     except Exception as e:
-        return f"❌ 保存失败: {str(e)}"
+        return f"❌ 上传失败: {str(e)}"
+
 
 def load_annotations(username):
-    """加载用户特定的标注结果（仅本地版本，Hub 上的可选扩展）"""
+    """从 Hugging Face Hub 加载用户特定的标注结果"""
     try:
-        annotation_file = get_user_annotation_file(username)
-        if os.path.exists(annotation_file):
-            with open(annotation_file, "r", encoding="utf-8") as f:
-                save_data = json.load(f)
-            annotation_results = {ann["task_id"]: ann["choice"] for ann in save_data["annotations"]}
-            return annotation_results
-        else:
-            return {}
-    except Exception as e:
+        filename = get_user_annotation_filename(username)
+        # 下载用户的标注文件
+        local_path = hf_hub_download(
+            repo_id=REPO_ID,
+            filename=filename,
+            repo_type="dataset",
+            token=HF_TOKEN,
+            force_download=True  # 确保拿到最新版本
+        )
+        with open(local_path, "r", encoding="utf-8") as f:
+            save_data = json.load(f)
+        annotation_results = {ann["task_id"]: ann["choice"] for ann in save_data.get("annotations", [])}
+        return annotation_results
+    except Exception:
+        # 用户还没有标注文件的情况
         return {}