Using datasets space

app.py

@@ -2,13 +2,27 @@ import sys
 import os
 import json
 import html
+import time
+import tempfile
 from datetime import datetime
 from pathlib import Path
 import gradio as gr
+try:
+    from huggingface_hub import HfApi, hf_hub_download
+    from huggingface_hub.utils import EntryNotFoundError
+except Exception:
+    HfApi = None
+    hf_hub_download = None
+    EntryNotFoundError = Exception
 
 DATA_DIR = Path(__file__).parent / "data"
 PAPERS_DIR = DATA_DIR / "papers"
 ANNOTATIONS_ROOT = DATA_DIR / "user_annotations"
+HF_DATASET_REPO = os.getenv("HF_DATASET_REPO", "ImanAndrea/citation_annotations")
+HF_ANNOTATIONS_PREFIX = os.getenv("HF_ANNOTATIONS_PREFIX", "annotations")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+HF_CACHE_DIR = Path(os.getenv("HF_CACHE_DIR", "/tmp/hf_cache"))
+HF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
 CATEGORIES = ["Unsupported claim", "Format", "Coherence", "Lacks synthesis"]
 CATEGORY_COLORS = {
     "Unsupported claim": "#ffb3b3",
@@ -30,6 +44,111 @@ LEGEND_HTML = """
 </div>
 """
 
+_REPO_FILES_CACHE = {"ts": 0.0, "files": []}
+
+
+def _use_hf_backend():
+    return bool(HF_TOKEN) and HF_DATASET_REPO and HfApi is not None and hf_hub_download is not None
+
+
+def _hf_api():
+    if not _use_hf_backend():
+        return None
+    return HfApi(token=HF_TOKEN)
+
+
+def _hf_repo_path(user_id, filename):
+    safe_user = _normalize_user_id(user_id)
+    if not safe_user:
+        return None
+    safe_file = _sanitize_filename(Path(filename).name)
+    return f"{HF_ANNOTATIONS_PREFIX}/{safe_user}/{safe_file}.json"
+
+
+def _hf_list_repo_files(ttl_seconds=30):
+    now = time.time()
+    if now - _REPO_FILES_CACHE["ts"] < ttl_seconds:
+        return _REPO_FILES_CACHE["files"]
+    api = _hf_api()
+    if not api:
+        return []
+    try:
+        files = api.list_repo_files(repo_id=HF_DATASET_REPO, repo_type="dataset")
+    except Exception:
+        files = []
+    _REPO_FILES_CACHE["ts"] = now
+    _REPO_FILES_CACHE["files"] = files
+    return files
+
+
+def _hf_download_annotations(repo_path):
+    if not repo_path:
+        return []
+    try:
+        local_path = hf_hub_download(
+            repo_id=HF_DATASET_REPO,
+            repo_type="dataset",
+            filename=repo_path,
+            token=HF_TOKEN,
+            local_dir=str(HF_CACHE_DIR),
+        )
+    except EntryNotFoundError:
+        return []
+    except Exception:
+        return []
+    try:
+        with open(local_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            return data if isinstance(data, list) else []
+    except Exception:
+        return []
+
+
+def _hf_upload_annotations(repo_path, records, commit_message):
+    api = _hf_api()
+    if not api:
+        raise RuntimeError("Hugging Face backend unavailable.")
+    HF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json", dir=str(HF_CACHE_DIR))
+    try:
+        with open(tmp_file.name, "w", encoding="utf-8") as f:
+            json.dump(records, f, ensure_ascii=False, indent=2)
+        api.upload_file(
+            path_or_fileobj=tmp_file.name,
+            path_in_repo=repo_path,
+            repo_id=HF_DATASET_REPO,
+            repo_type="dataset",
+            commit_message=commit_message,
+        )
+    finally:
+        try:
+            os.remove(tmp_file.name)
+        except Exception:
+            pass
+    _REPO_FILES_CACHE["ts"] = 0.0
+
+
+def _hf_delete_user_annotations(user_id):
+    api = _hf_api()
+    if not api:
+        raise RuntimeError("Hugging Face backend unavailable.")
+    safe_user = _normalize_user_id(user_id)
+    if not safe_user:
+        return
+    prefix = f"{HF_ANNOTATIONS_PREFIX}/{safe_user}/"
+    files = [f for f in _hf_list_repo_files() if f.startswith(prefix)]
+    for path in files:
+        try:
+            api.delete_file(
+                repo_id=HF_DATASET_REPO,
+                repo_type="dataset",
+                path_in_repo=path,
+                commit_message=f"Delete annotations for {safe_user}",
+            )
+        except Exception:
+            pass
+    _REPO_FILES_CACHE["ts"] = 0.0
+
 CATEGORY_JS = """
 (start, end) => {
   const findById = (id) => {
@@ -90,16 +209,24 @@ def list_papers():
 
 def list_user_ids():
     ids = set()
-    if ANNOTATIONS_ROOT.exists():
-        for path in ANNOTATIONS_ROOT.iterdir():
-            if path.is_dir():
-                ids.add(path.name)
-    # Legacy single-file annotations
-    base = Path(__file__).parent
-    for path in base.glob("annotations_*.json"):
-        stem = path.stem
-        if stem.startswith("annotations_"):
-            ids.add(stem[len("annotations_") :])
+    if _use_hf_backend():
+        for path in _hf_list_repo_files():
+            if not path.startswith(f"{HF_ANNOTATIONS_PREFIX}/"):
+                continue
+            parts = path.split("/")
+            if len(parts) >= 3 and parts[1]:
+                ids.add(parts[1])
+    else:
+        if ANNOTATIONS_ROOT.exists():
+            for path in ANNOTATIONS_ROOT.iterdir():
+                if path.is_dir():
+                    ids.add(path.name)
+        # Legacy single-file annotations
+        base = Path(__file__).parent
+        for path in base.glob("annotations_*.json"):
+            stem = path.stem
+            if stem.startswith("annotations_"):
+                ids.add(stem[len("annotations_") :])
     return sorted(ids)
 
 
@@ -206,6 +333,19 @@ def _annotations_path(user_id, filename):
 
 
 def _load_annotations(user_id, filename=None):
+    if _use_hf_backend():
+        if filename:
+            repo_path = _hf_repo_path(user_id, filename)
+            return _hf_download_annotations(repo_path)
+        safe_user = _normalize_user_id(user_id)
+        if not safe_user:
+            return []
+        prefix = f"{HF_ANNOTATIONS_PREFIX}/{safe_user}/"
+        records = []
+        for path in _hf_list_repo_files():
+            if path.startswith(prefix) and path.endswith(".json"):
+                records.extend(_hf_download_annotations(path))
+        return records
     user_dir = _user_dir(user_id)
     if not user_dir:
         return []
@@ -253,6 +393,16 @@ def _load_annotations(user_id, filename=None):
 
 
 def _save_annotations(user_id, filename, records):
+    if _use_hf_backend():
+        repo_path = _hf_repo_path(user_id, filename)
+        if not repo_path:
+            raise ValueError("Missing user ID or filename.")
+        _hf_upload_annotations(
+            repo_path,
+            records,
+            commit_message=f"Update annotations for {_normalize_user_id(user_id)}/{_sanitize_filename(filename)}",
+        )
+        return
     path = _annotations_path(user_id, filename)
     if not path:
         raise ValueError("Missing user ID or filename.")
@@ -308,6 +458,8 @@ def save_annotation(user_id, filename, start, end, category):
         _save_annotations(user_id, filename, records)
         content = _render_with_highlights(text, _annotations_for(user_id, filename))
         annotations_json = load_annotations_for(user_id, filename)
+        if _use_hf_backend():
+            return f"Saved annotation for {filename}.", content, annotations_json
         try:
             path = _annotations_path(user_id, filename)
             size = os.path.getsize(path)
@@ -360,6 +512,8 @@ def remove_annotation(user_id, filename, start, end):
         annotations_json = json.dumps(filtered, ensure_ascii=False, indent=2)
         if removed == 0:
             return "No overlapping highlights to remove.", content, annotations_json
+        if _use_hf_backend():
+            return f"Removed {removed} highlight(s).", content, annotations_json
         try:
             path = _annotations_path(user_id, filename)
             size = os.path.getsize(path)
@@ -378,19 +532,22 @@ def clear_all_annotations(user_id):
     if not user_id:
         return "Please enter your ID.", gr.update(), gr.update()
     try:
-        user_dir = _user_dir(user_id)
-        if user_dir and user_dir.exists():
-            for path in user_dir.glob("*.json"):
+        if _use_hf_backend():
+            _hf_delete_user_annotations(user_id)
+        else:
+            user_dir = _user_dir(user_id)
+            if user_dir and user_dir.exists():
+                for path in user_dir.glob("*.json"):
+                    try:
+                        path.unlink()
+                    except Exception:
+                        pass
+            legacy = Path(__file__).parent / f"annotations_{_normalize_user_id(user_id)}.json"
+            if legacy.exists():
                 try:
-                    path.unlink()
+                    legacy.unlink()
                 except Exception:
                     pass
-        legacy = Path(__file__).parent / f"annotations_{_normalize_user_id(user_id)}.json"
-        if legacy.exists():
-            try:
-                legacy.unlink()
-            except Exception:
-                pass
         return "Cleared all annotations.", "<em>No file selected.</em>", "[]"
     except Exception as e:
         return f"Error clearing annotations: {e}", gr.update(), gr.update()
@@ -409,9 +566,10 @@ def _add_user_id(new_id):
     safe = _normalize_user_id(new_id)
     if not safe:
         return gr.update(), "**Please enter a valid ID.**", gr.update()
-    ANNOTATIONS_ROOT.mkdir(parents=True, exist_ok=True)
-    user_dir = ANNOTATIONS_ROOT / safe
-    user_dir.mkdir(parents=True, exist_ok=True)
+    if not _use_hf_backend():
+        ANNOTATIONS_ROOT.mkdir(parents=True, exist_ok=True)
+        user_dir = ANNOTATIONS_ROOT / safe
+        user_dir.mkdir(parents=True, exist_ok=True)
     choices = list_user_ids()
     if safe not in choices:
         choices.append(safe)

requirements.txt

@@ -3,6 +3,7 @@ pandas
 numpy 
 transformers
 datasets 
+huggingface_hub
 torch 
 scikit-learn 
-nltk
+nltk