Spaces:
Running
Running
| 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 = 4) -> 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_obj): | |
| if file_obj is None: | |
| return "Please upload a FASTA file" | |
| # Read the file content | |
| try: | |
| # Handle both string and file object cases | |
| if isinstance(file_obj, str): | |
| text = file_obj | |
| else: | |
| text = file_obj.decode('utf-8') | |
| except Exception as e: | |
| return f"Error reading file: {str(e)}" | |
| # Load model and scaler | |
| try: | |
| 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() | |
| except Exception as e: | |
| return f"Error loading model: {str(e)}" | |
| # Get predictions | |
| results = [] | |
| try: | |
| 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) | |
| except Exception as e: | |
| return f"Error processing sequences: {str(e)}" | |
| return "\n\n".join(results) | |
| # Create the interface | |
| iface = gr.Interface( | |
| fn=predict, | |
| inputs=gr.File(label="Upload FASTA file", type="binary"), | |
| outputs=gr.Textbox(label="Results"), | |
| title="Virus Host Classifier" | |
| ) | |
| # Launch the interface | |
| if __name__ == "__main__": | |
| iface.launch() # Remove share=True for Hugging Face Spaces |