Update app.py

app.py

@@ -25,14 +25,27 @@ 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:
+    """Convert sequence to k-mer frequency vector"""
+    # Generate all possible k-mers
     kmers = [''.join(p) for p in product("ACGT", repeat=k)]
-    kmer_dict = {kmer: 0 for kmer in kmers}
+    kmer_dict = {km: i for i, km in enumerate(kmers)}
+
+    # Initialize vector
+    vec = np.zeros(len(kmers), dtype=np.float32)
+    
+    # Count k-mers
     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()))
+            vec[kmer_dict[kmer]] += 1
+
+    # Convert to frequencies
+    total_kmers = len(sequence) - k + 1
+    if total_kmers > 0:
+        vec = vec / total_kmers
+
+    return vec
 
 def parse_fasta(text):
     sequences = []
@@ -71,12 +84,19 @@ def predict(file_obj):
     # 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))
+        model = VirusClassifier(256).to(device)  # k=4 -> 4^4 = 256 features
+        
+        # Load model with explicit map_location
+        state_dict = torch.load('model.pt', map_location=device)
+        model.load_state_dict(state_dict)
+        
+        # Load scaler
         scaler = joblib.load('scaler.pkl')
+        
+        # Set model to evaluation mode
         model.eval()
     except Exception as e:
-        return f"Error loading model: {str(e)}"
+        return f"Error loading model: {str(e)}\nFull traceback: {str(e.__traceback__)}"
 
     # Get predictions
     results = []