Update app.py

app.py

@@ -9,34 +9,12 @@ import torch.nn.functional as F
 from Bio.Blast import NCBIWWW, NCBIXML
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel
 from huggingface_hub import hf_hub_download
+import spaces  # Safe to import at the top now!
 
 # ===================================
-# 1. LOAD AI MODELS (GLOBALLY CACHED)
+# 1. GLOBAL VARIABLES & CACHE
 # ===================================
-# Because "import spaces" is not called yet, PyTorch is safe from the Meta Device!
-print("Waking up the Genomic Oracle... Loading models safely into CPU RAM.\n")
-
-# A. Kadir's Gatekeeper 
-clf_coding = joblib.load("coding_classifier_universal.joblib")
-
-# B. Base DNABERT 
-tokenizer_base = AutoTokenizer.from_pretrained("DNABERT_Local", trust_remote_code=True)
-model_base = AutoModel.from_pretrained("DNABERT_Local", trust_remote_code=True, low_cpu_mem_usage=False)
-model_base.eval()
-
-# C. DNABERT-2 Promoter Model 
-tokenizer_promoter = AutoTokenizer.from_pretrained("llm_promoter_classifier_v2", trust_remote_code=True)
-model_promoter = AutoModelForSequenceClassification.from_pretrained("llm_promoter_classifier_v2", trust_remote_code=True, low_cpu_mem_usage=False)
-model_promoter.eval()
-
-# D. Multi-Feature LightGBM 
-lgbm_path = hf_hub_download(repo_id="Geonomic/Genomic-Oracle-Weights", filename="dnabert_lightgbm_model_feature_type.pkl")
-lightgbm_model = joblib.load(lgbm_path)
-
-# E. Custom ALiBi Lean/Obese BERT 
-tokenizer_pheno = AutoTokenizer.from_pretrained("Geonomic/Genomic-Oracle-Weights", trust_remote_code=True)
-model_pheno = AutoModelForSequenceClassification.from_pretrained("Geonomic/Genomic-Oracle-Weights", trust_remote_code=True, low_cpu_mem_usage=False)
-model_pheno.eval()
+print("Waking up the Genomic Oracle....\n")
 
 # Structural Feature Dictionary
 FEATURE_DICT = {
@@ -44,16 +22,49 @@ FEATURE_DICT = {
     3: "3' UTR", 4: "Promoter", 5: "Enhancer", 6: "lncRNA"
 }
 
+# This dictionary caches the models so they only load ONCE during the first click
+ORACLE_BRAIN = {}
+
 # ==============================================
 # 2. CORE INFERENCE ENGINE (ZeroGPU Accelerated)
 # ==============================================
-import spaces  # <--- THE MAGIC DELAY: We import ZeroGPU ONLY AFTER models are safely loaded!
-
-@spaces.GPU  # This tells Hugging Face to teleport the CPU models to the A100 GPU dynamically
+@spaces.GPU  # Teleports this specific math to the A100 GPU
 def run_deep_learning_cascade(dna_sequence):
     device = torch.device("cuda")  # Wakes up the A100 connection
     clean_seq = "".join(dna_sequence.split()).upper()
     
+    # --- JUST-IN-TIME CACHED LOADING ---
+    if "model_base" not in ORACLE_BRAIN:
+        print("🚀 Allocated GPU is Awake! Loading ALiBi models into VRAM...")
+        
+        ORACLE_BRAIN["clf_coding"] = joblib.load("coding_classifier_universal.joblib")
+        
+        ORACLE_BRAIN["tokenizer_base"] = AutoTokenizer.from_pretrained("DNABERT_Local", trust_remote_code=True)
+        ORACLE_BRAIN["model_base"] = AutoModel.from_pretrained("DNABERT_Local", trust_remote_code=True, low_cpu_mem_usage=False).to(device)
+        ORACLE_BRAIN["model_base"].eval()
+        
+        ORACLE_BRAIN["tokenizer_promoter"] = AutoTokenizer.from_pretrained("llm_promoter_classifier_v2", trust_remote_code=True)
+        ORACLE_BRAIN["model_promoter"] = AutoModelForSequenceClassification.from_pretrained("llm_promoter_classifier_v2", trust_remote_code=True, low_cpu_mem_usage=False).to(device)
+        ORACLE_BRAIN["model_promoter"].eval()
+        
+        lgbm_path = hf_hub_download(repo_id="Geonomic/Genomic-Oracle-Weights", filename="dnabert_lightgbm_model_feature_type.pkl")
+        ORACLE_BRAIN["lightgbm_model"] = joblib.load(lgbm_path)
+        
+        ORACLE_BRAIN["tokenizer_pheno"] = AutoTokenizer.from_pretrained("Geonomic/Genomic-Oracle-Weights", trust_remote_code=True)
+        ORACLE_BRAIN["model_pheno"] = AutoModelForSequenceClassification.from_pretrained("Geonomic/Genomic-Oracle-Weights", trust_remote_code=True, low_cpu_mem_usage=False).to(device)
+        ORACLE_BRAIN["model_pheno"].eval()
+        print("✅ All models successfully loaded into the Oracle Brain!")
+
+    # Retrieve from cache for instant inference
+    clf_coding = ORACLE_BRAIN["clf_coding"]
+    tokenizer_base = ORACLE_BRAIN["tokenizer_base"]
+    model_base = ORACLE_BRAIN["model_base"]
+    tokenizer_promoter = ORACLE_BRAIN["tokenizer_promoter"]
+    model_promoter = ORACLE_BRAIN["model_promoter"]
+    lightgbm_model = ORACLE_BRAIN["lightgbm_model"]
+    tokenizer_pheno = ORACLE_BRAIN["tokenizer_pheno"]
+    model_pheno = ORACLE_BRAIN["model_pheno"]
+
     # --- LEVEL 1: Base Embedding & Kadir's Gatekeeper ---
     inputs = tokenizer_base([clean_seq], return_tensors="pt", max_length=300, truncation=True, padding=True)
     inputs = {k: v.to(device) for k, v in inputs.items()}
@@ -67,7 +78,7 @@ def run_deep_learning_cascade(dna_sequence):
     p_coding = clf_coding.predict_proba(vector)[0][1]
     is_coding = p_coding >= 0.60
     
-    raw_scores = {"Gatekeeper (Gene Likelihood)": p_coding}
+    raw_scores = {"Gatekeeper (Gene Prospect)": p_coding}
 
     # --- LEVEL 2: LightGBM Structural Classification ---
     lgb_prediction = int(lightgbm_model.predict(vector)[0])
@@ -207,8 +218,7 @@ def gradio_inference(dna_sequence, run_mapping):
     return summary, "\n".join(stats_lines), context_output, ""
 
 # --- THE UI LAYOUT ---
-with gr.Blocks(theme=gr.themes.Soft(), title="🧬 The Genomic Oracle") as demo:
-    # The Integrated Landing Page
+with gr.Blocks(theme=gr.themes.Soft(), title="🧬 The Genomic Oracle 🧬") as demo:
     gr.Markdown(
         """
         # 🧬 The Genomic Oracle
@@ -240,7 +250,6 @@ with gr.Blocks(theme=gr.themes.Soft(), title="🧬 The Genomic Oracle") as demo:
 
     info_box = gr.Markdown("", elem_id="info_box")
 
-    # Connect the UI to the Python logic
     submit_btn.click(
         fn=lambda seq, map_flag: (
             *gradio_inference(seq, map_flag)[:3],