Update app.py

app.py

@@ -14,6 +14,7 @@ import pickle
 import torch
 from transformers import BertModel, BertTokenizer
 from model import PredictModel
+from huggingface_hub import snapshot_download
 
 # Import your local helper files
 # (These must be uploaded to the same folder as app.py)
@@ -25,9 +26,17 @@ SEEDS = [2023, 2024, 2025]
 TASK = 'PDC'
 
 # --- 3. LOAD PROTBERT (The "Brain" for Peptides) ---
-print("📦 Downloading/Loading ProtBERT Model... (This runs once)")
-tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
-bert_model = BertModel.from_pretrained("Rostlab/prot_bert")
+print("📦 Downloading ProtBERT Model safely...")
+
+# SMART DOWNLOAD: Only download the PyTorch files (skips the huge TensorFlow file)
+model_path = snapshot_download(
+    repo_id="Rostlab/prot_bert",
+    allow_patterns=["*.json", "pytorch_model.bin", "vocab.txt"], 
+    local_files_only=False
+)
+
+tokenizer = BertTokenizer.from_pretrained(model_path, do_lower_case=False)
+bert_model = BertModel.from_pretrained(model_path)
 
 # Move to CPU (Hugging Face Free Tier is CPU-only usually)
 bert_model.eval()