app.py

import gradio as gr
import torch
import joblib
import numpy as np
from itertools import product
import torch.nn as nn

class VirusClassifier(nn.Module):
    def __init__(self, input_shape: int):
        super(VirusClassifier, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_shape, 64),
            nn.GELU(),
            nn.BatchNorm1d(64),
            nn.Dropout(0.3),
            nn.Linear(64, 32),
            nn.GELU(),
            nn.BatchNorm1d(32),
            nn.Dropout(0.3),
            nn.Linear(32, 32),
            nn.GELU(),
            nn.Linear(32, 2)
        )

    def forward(self, x):
        return self.network(x)

def sequence_to_kmer_vector(sequence: str, k: int = 6) -> np.ndarray:
    kmers = [''.join(p) for p in product("ACGT", repeat=k)]
    kmer_dict = {kmer: 0 for kmer in kmers}
    
    for i in range(len(sequence) - k + 1):
        kmer = sequence[i:i+k]
        if kmer in kmer_dict:
            kmer_dict[kmer] += 1
            
    return np.array(list(kmer_dict.values()))

def parse_fasta(text):
    sequences = []
    current_header = None
    current_sequence = []
    
    for line in text.split('\n'):
        line = line.strip()
        if not line:
            continue
        if line.startswith('>'):
            if current_header:
                sequences.append((current_header, ''.join(current_sequence)))
            current_header = line[1:]
            current_sequence = []
        else:
            current_sequence.append(line.upper())
            
    if current_header:
        sequences.append((current_header, ''.join(current_sequence)))
        
    return sequences

def predict(file):
    if not file:
        return "Please upload a FASTA file"
    
    # Read the file
    text = file.decode('utf-8')
    
    # Load model and scaler
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = VirusClassifier(4096).to(device)
    model.load_state_dict(torch.load('model.pt', map_location=device))
    scaler = joblib.load('scaler.pkl')
    model.eval()
    
    # Get predictions
    results = []
    sequences = parse_fasta(text)
    
    for header, seq in sequences:
        # Get k-mer vector
        kmer_vector = sequence_to_kmer_vector(seq)
        kmer_vector = scaler.transform(kmer_vector.reshape(1, -1))
        
        # Predict
        with torch.no_grad():
            output = model(torch.FloatTensor(kmer_vector).to(device))
            probs = torch.softmax(output, dim=1)
        
        # Format results
        pred_class = 1 if probs[0][1] > probs[0][0] else 0
        pred_label = 'human' if pred_class == 1 else 'non-human'
        
        result = f"""
Sequence: {header}
Prediction: {pred_label}
Confidence: {float(max(probs[0])):0.4f}
Human probability: {float(probs[0][1]):0.4f}
Non-human probability: {float(probs[0][0]):0.4f}
"""
        results.append(result)
    
    return "\n".join(results)

# Create the interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.File(label="Upload FASTA file"),
    outputs=gr.Textbox(label="Results"),
    title="Virus Host Classifier"
)

# Launch
iface.launch()