Update app.py

app.py

@@ -2,22 +2,25 @@
 import os, shutil, subprocess
 from huggingface_hub import scan_cache_dir, snapshot_download
 
-# 1) Put ALL caches in /data so they’re manageable & persistent
-os.makedirs("/data/.cache", exist_ok=True)
+# Put caches in /data and make sure dirs exist
+os.makedirs("/data/.cache/huggingface/hub", exist_ok=True)
+os.makedirs("/data/snapshots", exist_ok=True)
+
 os.environ.setdefault("XDG_CACHE_HOME", "/data/.cache")
 os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
 os.environ.setdefault("HF_HUB_CACHE", "/data/.cache/huggingface/hub")
-os.environ.setdefault("TRANSFORMERS_CACHE", "/data/.cache/huggingface/transformers")
-os.environ.setdefault("DATASETS_CACHE", "/data/.cache/huggingface/datasets")
+# Avoid TRANSFORMERS_CACHE deprecation; HF_HOME is enough.
+# os.environ.setdefault("TRANSFORMERS_CACHE", "/data/.cache/huggingface/transformers")
 
-# 2) Prune old HF cache revisions (keeps current blobs, deletes stale revs)
+# Prune old HF cache revisions (safe if empty; now the dir exists)
 try:
     cache = scan_cache_dir(os.environ["HF_HUB_CACHE"])
-    cache.delete_revisions([rev for rev in cache.revisions])
+    if cache.revisions:
+        cache.delete_revisions([rev for rev in cache.revisions])
 except Exception as e:
     print(f"[cache prune] skipped: {e}")
 
-# (Optional) light guard: trim pip wheel cache
+# Light pip cache cleanup
 try:
     subprocess.run(["pip", "cache", "purge"], check=False)
 except Exception:
@@ -29,48 +32,54 @@ import sys
 import pandas as pd
 from transformers import AutoTokenizer, AutoModel, AutoConfig
 
-# If you want fully reproducible rebuilds, set these in Space → Settings → Variables
-# (or leave blank to use latest)
+# Optional: pin commits via Space Variables
 MODEL_ID = "ChatterjeeLab/MetaLATTE"
 TOKENIZER_ID = "facebook/esm2_t33_650M_UR50D"
 MODEL_REV = os.getenv("MODEL_REV", "")         # e.g. "a1b2c3d"
 TOKENIZER_REV = os.getenv("TOKENIZER_REV", "") # e.g. "9f8e7d6"
 
-# Prefer downloading *exactly* what you need to /data and load locally.
-# This avoids multiple revision copies over time.
-def maybe_snapshot(repo_id, revision, allow_patterns):
-    kw = dict(repo_id=repo_id, local_dir=None, ignore_regex=None)
-    if revision:
-        kw["revision"] = revision
-    # Download to HF cache in /data; return the resolved local dir
-    return snapshot_download(allow_patterns=allow_patterns, **kw)
-
-# Download tokenizer files only (small)
-esm_local = maybe_snapshot(
-    TOKENIZER_ID, TOKENIZER_REV,
+def snapshot_to(local_name, repo_id, revision, allow_patterns):
+    """Download only needed files into a concrete folder under /data/snapshots."""
+    local_dir = f"/data/snapshots/{local_name}"
+    os.makedirs(local_dir, exist_ok=True)
+    # IMPORTANT: no ignore_regex; use ignore_patterns if needed
+    return snapshot_download(
+        repo_id=repo_id,
+        revision=revision if revision else None,
+        allow_patterns=allow_patterns,
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,   # copy files into local_dir; easier to manage size
+    )
+
+# Tokenizer (small set of files)
+esm_local = snapshot_to(
+    "esm2_tokenizer",
+    TOKENIZER_ID,
+    TOKENIZER_REV,
     allow_patterns=[
         "tokenizer.json","tokenizer_config.json","vocab.*","merges.*",
         "special_tokens_map.json","*.model","tokenizer*.txt","spiece.*","*.tiktoken"
-    ]
+    ],
 )
 
-# Download MetaLATTE (weights + config only)
-metalatte_local = maybe_snapshot(
-    MODEL_ID, MODEL_REV,
-    allow_patterns=["*.json","*.safetensors","*.bin","*.model","*.txt"]  # keep it tight
+# MetaLATTE model (weights + config only)
+metalatte_local = snapshot_to(
+    "metalatte_model",
+    MODEL_ID,
+    MODEL_REV,
+    allow_patterns=["*.json","*.safetensors","*.bin","*.model","*.txt"],
 )
 
-# Add the current directory to the system path for your custom code
+# Your local package
 metalatte_path = '.'
 sys.path.insert(0, metalatte_path)
 
-# Import the custom configuration and model
 from configuration import MetaLATTEConfig
 from modeling_metalatte import MultitaskProteinModel
 AutoConfig.register("metalatte", MetaLATTEConfig)
 AutoModel.register(MetaLATTEConfig, MultitaskProteinModel)
 
-# Load from the local snapshot dirs (avoids re-downloading on rebuilds)
+# Load from the downloaded dirs (no network, no extra cache growth)
 tokenizer = AutoTokenizer.from_pretrained(esm_local, local_files_only=True)
 config = AutoConfig.from_pretrained(metalatte_local, local_files_only=True)
 model = AutoModel.from_pretrained(metalatte_local, config=config, local_files_only=True)