|
|
| """
|
| 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")
|
|
|
|
|
| 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"]
|
|
|
|
|
| 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)
|
|
|
| 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()
|
|
|
|
|
| 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)
|
|
|
| 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() |
