hiyata commited on
Commit
5263bd3
·
verified ·
1 Parent(s): be92930

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +116 -0
app.py ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import joblib
4
+ import numpy as np
5
+ from itertools import product
6
+ from typing import Dict
7
+ import torch.nn as nn
8
+
9
+ class VirusClassifier(nn.Module):
10
+ def __init__(self, input_shape: int):
11
+ super(VirusClassifier, self).__init__()
12
+ self.network = nn.Sequential(
13
+ nn.Linear(input_shape, 64),
14
+ nn.GELU(),
15
+ nn.BatchNorm1d(64),
16
+ nn.Dropout(0.3),
17
+ nn.Linear(64, 32),
18
+ nn.GELU(),
19
+ nn.BatchNorm1d(32),
20
+ nn.Dropout(0.3),
21
+ nn.Linear(32, 32),
22
+ nn.GELU(),
23
+ nn.Linear(32, 2)
24
+ )
25
+
26
+ def forward(self, x):
27
+ return self.network(x)
28
+
29
+ def sequence_to_kmer_vector(sequence: str, k: int = 6) -> np.ndarray:
30
+ """Convert sequence to k-mer frequency vector"""
31
+ kmers = [''.join(p) for p in product("ACGT", repeat=k)]
32
+ kmer_dict = {kmer: 0 for kmer in kmers}
33
+
34
+ for i in range(len(sequence) - k + 1):
35
+ kmer = sequence[i:i+k]
36
+ if kmer in kmer_dict: # only count valid kmers
37
+ kmer_dict[kmer] += 1
38
+
39
+ return np.array(list(kmer_dict.values()))
40
+
41
+ def parse_fasta(fasta_content: str):
42
+ """Parse FASTA format string"""
43
+ sequences = []
44
+ current_header = None
45
+ current_sequence = []
46
+
47
+ for line in fasta_content.split('\n'):
48
+ line = line.strip()
49
+ if not line:
50
+ continue
51
+ if line.startswith('>'):
52
+ if current_header is not None:
53
+ sequences.append((current_header, ''.join(current_sequence)))
54
+ current_header = line[1:]
55
+ current_sequence = []
56
+ else:
57
+ current_sequence.append(line.upper())
58
+
59
+ if current_header is not None:
60
+ sequences.append((current_header, ''.join(current_sequence)))
61
+
62
+ return sequences
63
+
64
+ def predict_sequence(fasta_content: str) -> str:
65
+ """Process FASTA input and return formatted predictions"""
66
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
67
+ k = 6
68
+
69
+ # Load model and scaler
70
+ model = VirusClassifier(4096).to(device) # 4096 = 4^6 for 6-mers
71
+ model.load_state_dict(torch.load('model.pt', map_location=device))
72
+ scaler = joblib.load('scaler.pkl')
73
+ model.eval()
74
+
75
+ # Process sequences
76
+ sequences = parse_fasta(fasta_content)
77
+ results = []
78
+
79
+ for header, seq in sequences:
80
+ # Convert sequence to k-mer vector
81
+ kmer_vector = sequence_to_kmer_vector(seq, k)
82
+ kmer_vector = scaler.transform(kmer_vector.reshape(1, -1))
83
+
84
+ # Get prediction
85
+ with torch.no_grad():
86
+ output = model(torch.FloatTensor(kmer_vector).to(device))
87
+ probs = torch.softmax(output, dim=1)
88
+
89
+ # Format result
90
+ pred_class = 1 if probs[0][1] > probs[0][0] else 0
91
+ pred_label = 'human' if pred_class == 1 else 'non-human'
92
+
93
+ result = f"""
94
+ Sequence: {header}
95
+ Prediction: {pred_label}
96
+ Confidence: {float(max(probs[0])):0.4f}
97
+ Human probability: {float(probs[0][1]):0.4f}
98
+ Non-human probability: {float(probs[0][0]):0.4f}
99
+ """
100
+ results.append(result)
101
+
102
+ return "\n".join(results)
103
+
104
+ # Create Gradio interface
105
+ iface = gr.Interface(
106
+ fn=predict_sequence,
107
+ inputs=gr.File(label="Upload FASTA file", file_types=[".fasta", ".fa", ".txt"]),
108
+ outputs=gr.Textbox(label="Prediction Results"),
109
+ title="Virus Host Classifier",
110
+ description="Upload a FASTA file to predict whether a virus sequence is likely to infect human or non-human hosts.",
111
+ examples=[["example.fasta"]],
112
+ cache_examples=True
113
+ )
114
+
115
+ # Launch the interface
116
+ iface.launch()