import spaces # MUST BE LINE 1. Fixes the "CUDA Initialized" error!
import gradio as gr
import os
import re
import requests
import torch
import joblib
import numpy as np
import torch.nn.functional as F
from Bio.Blast import NCBIWWW, NCBIXML
# include HF native imports for the Phenotype model
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel, BertTokenizer, BertForSequenceClassification, AutoConfig
from huggingface_hub import hf_hub_download
# ===================================
# 1. LOAD AI MODELS (GLOBALLY CACHED)
# ===================================
print("Waking up the Genomic Oracle...\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, _fast_init=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, _fast_init=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_v2.pkl")
raw_lgbm = joblib.load(lgbm_path)
# If it's a dictionary, print the keys to the log and try to extract the model
if isinstance(raw_lgbm, dict):
print(f" DEBUG: LightGBM Dictionary Keys: {raw_lgbm.keys()}")
# We will try the most common names for saved models
if "model" in raw_lgbm:
lightgbm_model = raw_lgbm["model"]
elif "classifier" in raw_lgbm:
lightgbm_model = raw_lgbm["classifier"]
else:
# Fallback: just grab the very first thing in the dictionary
first_key = list(raw_lgbm.keys())[0]
lightgbm_model = raw_lgbm[first_key]
else:
lightgbm_model = raw_lgbm
# E. Custom Lean/Obese Phenotype BERT (Forced Native Architecture via Colab Fix)
tokenizer_pheno = BertTokenizer.from_pretrained("Geonomic/Genomic-Oracle-Weights", do_lower_case=False)
config_pheno = AutoConfig.from_pretrained("Geonomic/Genomic-Oracle-Weights", trust_remote_code=True)
model_pheno = BertForSequenceClassification.from_pretrained("Geonomic/Genomic-Oracle-Weights", config=config_pheno, _fast_init=False)
model_pheno.eval()
FEATURE_DICT = {
0: "Gene/Transcript (Coding/mRNA)",
1: "Regulatory Region (Promoter/Enhancer/Silencer)",
2: "Long Non-Coding RNA (lncRNA)",
3: "Small/Transfer RNA (snRNA/miRNA/tRNA)",
4: "Repeat Region / Mobile Genetic Element",
5: "Pseudogene"
}
# ==============================================
# 2. CORE INFERENCE ENGINE (ZeroGPU Accelerated)
# ==============================================
@spaces.GPU
def run_deep_learning_cascade(dna_sequence):
device = torch.device("cuda")
# THE FINAL KEY: Teleport the CPU-locked models into the A100 GPU!
model_base.to(device)
model_promoter.to(device)
model_pheno.to(device)
clean_seq = "".join(dna_sequence.split()).upper()
# --- 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()}
with torch.no_grad():
out_base = model_base(**inputs)
mask = inputs["attention_mask"].unsqueeze(-1)
embedding = (out_base[0] * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
vector = embedding.float().cpu().numpy()
p_coding = clf_coding.predict_proba(vector)[0][1]
is_coding = p_coding >= 0.60
raw_scores = {"Protein-Coding Probability": p_coding}
# --- LEVEL 2: LightGBM Structural Classification ---
lgb_prediction = int(lightgbm_model.predict(vector)[0])
structural_feature = FEATURE_DICT.get(lgb_prediction, "Unknown Region")
# raw_scores["Predicted Structure"] = structural_feature
# THE CONTRADICTION RESOLVER
# The fine-tuned LightGBM overrides any contradiction during class assignment
if lgb_prediction == 0:
is_coding = True
# If LightGBM had to overrule, we boost the base confidence to match its high AUROC accuracy
confidence = p_coding if p_coding >= 0.50 else 0.85
else:
is_coding = False
confidence = (1 - p_coding) if p_coding < 0.50 else 0.85
# --- LEVEL 3: The Deep Learning Branching Logic ---
summary_dict = {
"Final Classification": "GENE" if is_coding else "NON-CODING",
"Feature": structural_feature
}
# BRANCH A: Phenotype Analysis (Triggered if Coding AND is CDS/Exon)
if is_coding and lgb_prediction == 0:
kmers = [clean_seq[i:i+5] for i in range(len(clean_seq) - 4)]
spaced_kmers = " ".join(kmers)
inputs_pheno = tokenizer_pheno(spaced_kmers, return_tensors="pt", max_length=512, truncation=True).to(device)
with torch.no_grad():
outputs = model_pheno(**inputs_pheno)
probs = F.softmax(outputs.logits, dim=-1)
prob_obese, prob_lean = probs[0][0].item(), probs[0][1].item()
phenotype = "Obesity-Associated" if prob_obese > prob_lean else "Lean-Associated"
summary_dict["Phenotype"] = phenotype
raw_scores["Phenotype (Obese)"] = prob_obese
raw_scores["Phenotype (Lean)"] = prob_lean
# BRANCH B: Promoter Validation (Triggered if Non-Coding AND is Promoter/Enhancer)
elif not is_coding and lgb_prediction == 1:
inputs_promo = tokenizer_promoter([clean_seq], return_tensors="pt", max_length=300, truncation=True, padding=True).to(device)
with torch.no_grad():
outputs = model_promoter(**inputs_promo)
probs = F.softmax(outputs.logits, dim=-1)
p_promoter = probs[0][0].item()
validation = "High Confidence Regulatory Element" if p_promoter >= 0.50 else "Weak Regulatory Signal"
summary_dict["Validation"] = validation
raw_scores["Promoter Signal"] = p_promoter
return summary_dict, confidence, raw_scores
# ===================================
# 3. SPATIAL MAPPING (NCBI / ENSEMBL)
# ===================================
def get_genomic_context(sequence, is_coding):
feature_type = "CODING" if is_coding else "PROMOTER"
try:
# ask blast for 5 hits instead of 1 so we can hunt for the true chromosome
result_handle = NCBIWWW.qblast(
"blastn",
"nt",
sequence,
entrez_query="Homo sapiens[Organism] AND biomol_genomic[PROP]",
hitlist_size=50
)
blast_record = NCBIXML.read(result_handle)
except Exception as e:
return {"error": f"BLAST Connection Error: {e}"}
if not blast_record.alignments:
return {"error": "No human genome match found for this sequence."}
# Loop through the top hits and grab the first one that is an actual Chromosome
alignment = blast_record.alignments[0] # Default to the top hit
chrom = None
for aln in blast_record.alignments:
chrom_match = re.search(r"chromosome\s([0-9XYMT]+)", aln.title, re.IGNORECASE)
if chrom_match:
alignment = aln
chrom = chrom_match.group(1)
break # We found the chromosome, stop searching!
hsp = alignment.hsps[0]
location_string = f"Chromosome {chrom}" if chrom else f"Accession {alignment.accession}"
start, end = hsp.sbjct_start, hsp.sbjct_end
is_forward = (start < end)
strand_txt = "Forward (+)" if is_forward else "Reverse (-)"
if chrom is None:
return {"location": location_string, "start": start, "end": end, "strand": strand_txt, "metadata": "BLAST returned a localized record without a chromosome. Ensembl mapping skipped."}
search_start = min(start, end) if feature_type == "CODING" else (end if is_forward else max(1, end - 15000))
search_end = max(start, end) if feature_type == "CODING" else (end + 15000 if is_forward else end)
try:
response = requests.get(
f"https://rest.ensembl.org/overlap/region/human/{chrom}:{search_start}-{search_end}?feature=gene",
headers={"Accept": "application/json"}
)
response.raise_for_status()
genes = response.json()
except Exception as e:
return {"location": location_string, "start": start, "end": end, "strand": strand_txt, "metadata": f"Ensembl mapping unavailable: {e}"}
if not genes:
gene_desc = "No annotated genes found in this specific region."
else:
if feature_type == "PROMOTER":
genes.sort(key=lambda x: min(abs(x['start'] - end), abs(x['end'] - end)))
top_gene = genes[0]
name = top_gene.get('external_name', 'Unknown')
biotype = top_gene.get('biotype', 'Unknown').replace('_', ' ').title()
desc = top_gene.get('description', 'No description available.').split(' [')[0]
gene_desc = f"Matches Gene: {name} | Type: {biotype} | Function: {desc}" if feature_type == "CODING" else f"Regulates Downstream Gene: {name} | Type: {biotype} | Function: {desc}"
return {"location": location_string, "start": start, "end": end, "strand": strand_txt, "metadata": gene_desc}
# ==============================
# 4. GRADIO INTERFACE (FRONTEND)
# ==============================
def gradio_inference(dna_sequence, run_mapping):
if len(dna_sequence.strip()) < 10:
error_html = f"""
❌ Classification Summary
Sequence too short! Please enter at least 10 base pairs.
"""
# Updated the placeholder text here as well!
yield (error_html, "📊 Internal Pipeline Statistics
Results will appear here...
Internal Pipeline Statistics
"""
for key, val in raw_scores.items():
if isinstance(val, float):
# Added color: #10b981; to make all percentage scores GREEN
stats_html += f"{key}: {val:.2%}
"
else:
stats_html += f"{key}: {val}
"
stats_html += "
"
# Build Custom HTML for Summary
summary_html = f"""
Classification Summary (Deep Scan Complete)
Final Classification: {summary_dict.get('Final Classification')}
Feature: {summary_dict.get('Feature')}
"""
if "Phenotype" in summary_dict:
summary_html += f"Phenotype: {summary_dict['Phenotype']}
"
if "Validation" in summary_dict:
summary_html += f"Validation: {summary_dict['Validation']}
"
# Confidence score remains green
summary_html += f"Confidence Score: {conf:.2%}
"
if not run_mapping:
yield (summary_html, stats_html, "Spatial mapping skipped (Enable NCBI query to run).", "")
return
# Pushes AI results instantly while showing a loading message for BLAST!
yield (summary_html, stats_html, "⏳ Querying NCBI BLAST... (This takes 1-3 minutes. Please wait.)", "")
# Run Context Mapping (CPU / Network)
is_coding = summary_dict.get("Final Classification") == "GENE"
context = get_genomic_context(dna_sequence, is_coding)
if "error" in context:
context_output = f"❌ Mapping failed: {context['error']}"
elif "location" in context:
context_lines = [
f"Location:\t{context['location']}",
f"Strand:\t{context['strand']}",
f"Coordinates: {context['start']:,} – {context['end']:,}",
f"Notes:\t{context['metadata']}"
]
context_output = "\n".join(context_lines)
else:
context_output = "⚠️ Could not map sequence."
# Pushes the finished BLAST results!
yield (summary_html, stats_html, context_output, "")
# --- CUSTOM CSS ---
custom_css = """
#scan_btn {
background-color: #0d6efd !important; /* Deep Blue */
color: white !important;
border: none !important;
transition: 0.3s ease;
}
#scan_btn:hover {
background-color: #dc3545 !important; /* Striking Red */
}
"""
# --- THE UI LAYOUT ---
with gr.Blocks(theme=gr.themes.Soft(), title="🧬 The Genomic Oracle 🧬", css=custom_css) as demo:
# 1. The Custom HTML Title
gr.HTML(
"""
🧬 The Genomic Oracle 🧬
University of Maryland Global Campus | Bioinformatics Capstone
"""
)
# 2. The Standard Markdown Text
gr.Markdown(
"""
Welcome to the official interface for **The Genomic Oracle**, a cascaded machine learning pipeline designed for high-precision DNA sequence classification.
### The 4-Stage Cascading Architecture
1. **The Gatekeeper:** Logistic Regression model rapidly screens native k-mer vectors to identify protein-coding vs. non-coding potential.
2. **Structural Mapper:** LightGBM model classifies the sequence into 1 of 6 structural features (e.g., lncRNAs, Enhancers, mobile elements).
3. **Phenotype Prediction:** Sequences flagged as Coding are passed through a custom ALiBi BERT transformer to predict specific traits.
4. **Regulatory Validation:** Sequences flagged as Promoters or Enhancers are routed to a DNABERT-2 spatial attention neural network.
---
Created by: Kadir Galindo, Duncan Hall, Rebecca Mellinger & George Paccione\n
"""
)
with gr.Row():
with gr.Column(scale=1):
dna_input = gr.Textbox(label="Enter DNA Sequence", placeholder="e.g., ATGCGATCGATCGATCG...", lines=10)
run_mapping_cb = gr.Checkbox(value=False, label="Query NCBI BLAST for spatial mapping (Takes 1–3 mins)")
submit_btn = gr.Button("🚀 Initialize Deep Scan", elem_id="scan_btn")
with gr.Column(scale=1):
# Create default placeholder HTML so the boxes are visible on startup!
default_summary = """
Classification Summary
Results will appear here...
"""
default_stats = """
Internal Pipeline Statistics
Results will appear here...
"""
# Pass the default HTML into the components
output_summary = gr.HTML(value=default_summary)
stats_panel = gr.HTML(value=default_stats)
mapping_section = gr.Accordion("Genomic Context (BLAST/Ensembl)", open=False)
with mapping_section:
context_output = gr.Textbox(label="Mapping Results", lines=5, placeholder="Results will appear here...")
info_box = gr.Markdown("", elem_id="info_box")
# NEW UX FEATURE: Auto-open the accordion when the user checks the BLAST box!
run_mapping_cb.change(
fn=lambda is_checked: gr.Accordion(open=is_checked),
inputs=[run_mapping_cb],
outputs=[mapping_section]
)
# Pinned directly to the generator function to allow live-streaming!
submit_btn.click(
fn=gradio_inference,
inputs=[dna_input, run_mapping_cb],
outputs=[output_summary, stats_panel, context_output, info_box]
)
demo.launch()