Update app.py

app.py

@@ -25,15 +25,13 @@ class VirusClassifier(nn.Module):
     def forward(self, x):
         return self.network(x)
 
-def sequence_to_kmer_vector(sequence: str, k: int = 6) -> np.ndarray:
+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):
@@ -52,10 +50,8 @@ def parse_fasta(text):
             current_sequence = []
         else:
             current_sequence.append(line.upper())
-            
     if current_header:
         sequences.append((current_header, ''.join(current_sequence)))
-        
     return sequences
 
 def predict(file_obj):
@@ -63,51 +59,61 @@ def predict(file_obj):
         return "Please upload a FASTA file"
     
     # Read the file content
-    text = file_obj.read().decode()
-    
+    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
-    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()
-    
+    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 = []
-    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}
+    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)
-    
-    return "\n".join(results)
+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"),
+    inputs=gr.File(label="Upload FASTA file", type="binary"),
     outputs=gr.Textbox(label="Results"),
     title="Virus Host Classifier"
 )
 
-# Launch with public link
-iface.launch(share=True)
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch()  # Remove share=True for Hugging Face Spaces