Upload predict.py

predict.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""
+Viro-BERT: Standalone CLI Inference Script for Viral Pathogen Classification
+Compatible with Python 3.8+, PyTorch, and R (via reticulate or system calls).
+"""
+
+import os
+import sys
+import json
+import argparse
+import torch
+import torch.nn as nn
+from Bio import SeqIO
+from transformers import DistilBertConfig, DistilBertModel
+
+
+class BiologicalDistilBERT(nn.Module):
+    def __init__(self, vocab_size, num_classes):
+        super(BiologicalDistilBERT, self).__init__()
+        config = DistilBertConfig(
+            vocab_size=vocab_size,
+            max_position_embeddings=512,
+            n_layers=6,
+            n_heads=12,
+            dim=768,
+            dropout=0.1,
+            output_hidden_states=False
+        )
+        self.bert = DistilBertModel(config)
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 768),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(768, num_classes)
+        )
+
+    def forward(self, input_ids):
+        attention_mask = (input_ids != 0).long()
+        output = self.bert(input_ids, attention_mask=attention_mask)[0]
+        cls_output = output[:, 0, :]
+        return self.classifier(cls_output)
+
+
+def tokenize_sequence(seq, vocab, k=3, stride=3):
+    tokens = [vocab.get("[CLS]", 2)]
+    for i in range(0, len(seq) - k + 1, stride):
+        kmer = seq[i : i + k]
+        tokens.append(vocab.get(kmer, vocab.get("[UNK]", 1)))
+    tokens.append(vocab.get("[SEP]", 3))
+    return tokens
+
+
+def create_chunks(tokens, chunk_len=512, overlap=50):
+    stride = chunk_len - overlap
+    chunks = []
+    num_tokens = len(tokens)
+
+    if num_tokens <= chunk_len:
+        padded = tokens + [0] * (chunk_len - num_tokens)
+        chunks.append(padded)
+    else:
+        for i in range(0, num_tokens, stride):
+            chunk = tokens[i : i + chunk_len]
+            if len(chunk) > 10:
+                if len(chunk) < chunk_len:
+                    chunk = chunk + [0] * (chunk_len - len(chunk))
+                chunks.append(chunk)
+
+    return chunks
+
+
+def classify_fasta(fasta_path, weights_path, config_path, device=None, top_k=3):
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Load configuration
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Configuration file not found: {config_path}")
+    with open(config_path, "r", encoding="utf-8") as f:
+        config_data = json.load(f)
+
+    vocab = config_data["vocab"]
+    id_to_label = {int(k): v for k, v in config_data["id_to_label"].items()}
+    num_classes = config_data["num_classes"]
+
+    # Initialize and load model
+    model = BiologicalDistilBERT(vocab_size=len(vocab), num_classes=num_classes)
+    if not os.path.exists(weights_path):
+        raise FileNotFoundError(f"Weights file not found: {weights_path}")
+
+    state_dict = torch.load(weights_path, map_location=device)
+    # Clean keys if saved with torch.compile prefix
+    cleaned_state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
+    model.load_state_dict(cleaned_state_dict)
+    model.to(device)
+    model.eval()
+
+    # Parse FASTA sequences
+    records = list(SeqIO.parse(fasta_path, "fasta"))
+    if not records:
+        raise ValueError(f"No valid FASTA records found in: {fasta_path}")
+
+    print("=" * 65)
+    print(f"Viro-BERT Classification | Device: {device} | Sequences: {len(records)}")
+    print("=" * 65)
+
+    for record in records:
+        seq = str(record.seq).upper()
+        tokens = tokenize_sequence(seq, vocab, k=config_data.get("k_mer", 3), stride=config_data.get("stride", 3))
+        chunks = create_chunks(tokens, chunk_len=config_data.get("chunk_len", 512), overlap=config_data.get("overlap", 50))
+
+        if not chunks:
+            print(f"[!] Sequence {record.id} too short to extract tokens. Skipping.")
+            continue
+
+        input_tensor = torch.tensor(chunks, dtype=torch.long).to(device)
+
+        use_amp = (device.type == "cuda")
+        with torch.no_grad():
+            if use_amp:
+                with torch.autocast(device_type="cuda", dtype=torch.float16):
+                    logits = model(input_tensor)
+            else:
+                logits = model(input_tensor)
+
+            probs = torch.softmax(logits, dim=1)
+            # Late Fusion Consensus (Mean Pooling over chunks)
+            consensus_probs = torch.mean(probs, dim=0).cpu()
+
+        top_probs, top_indices = torch.topk(consensus_probs, k=min(top_k, num_classes))
+        winner_idx = top_indices[0].item()
+        winner_label = id_to_label[winner_idx]
+        winner_conf = top_probs[0].item() * 100
+
+        print(f"\nSequence ID   : {record.id}")
+        print(f"Length (bp)   : {len(seq)} bp ({len(chunks)} chunks processed)")
+        print(f"Prediction    : {winner_label} ({winner_conf:.2f}% confidence)")
+        print("Top Candidates:")
+        for rank, (p, idx) in enumerate(zip(top_probs, top_indices), start=1):
+            print(f"  {rank}. {id_to_label[idx.item()]:<10} : {p.item()*100:.2f}%")
+        print("-" * 65)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Viro-BERT: Whole-Genome Attention for Viral Pathogen Classification",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument("--fasta", type=str, required=True, help="Path to input FASTA file containing viral sequences")
+    parser.add_argument("--weights", type=str, default="sota_chunk_model.pth", help="Path to model weights (.pth)")
+    parser.add_argument("--config", type=str, default="config.json", help="Path to config.json containing vocab & label maps")
+    parser.add_argument("--top_k", type=int, default=3, help="Number of top candidates to report")
+    parser.add_argument("--cpu", action="store_true", help="Force CPU inference even if CUDA is available")
+
+    args = parser.parse_args()
+
+    device = torch.device("cpu") if args.cpu else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    classify_fasta(args.fasta, args.weights, args.config, device=device, top_k=args.top_k)
+
+
+if __name__ == "__main__":
+    main()