test

app.py

@@ -14,6 +14,9 @@ from PIL import Image
 from src.utils.get_features import get_img_api
 import joblib 
 
+
+device = torch.device('cpu')
+
 # Path to the dataset
 data_path = 'src/data/subset_dataset.csv'
 device = torch.device('cpu')
@@ -28,10 +31,14 @@ simple_transform = transforms.Compose([
 # Load the model
 def load_model(model_path, device='cpu'):
     """Loads the model from a joblib file and moves it to the specified device."""
-    model = joblib.load(model_path)
-    # If the model contains PyTorch tensors, move them to the specified device
+    # Use torch.load with map_location to ensure CPU compatibility
+    with open(model_path, 'rb') as f:
+        model = torch.load(f, map_location=device)
+    
+    # If the model is a PyTorch module, move it to the specified device and set to eval mode
     if isinstance(model, torch.nn.Module):
         model = model.to(device)
+        model.eval()
     return model
 
 # Get prediction
@@ -43,10 +50,11 @@ def get_prediction(model, padded_sequences, img_x, device='cpu'):
     padded_sequences, img_x = padded_sequences.to(device), img_x.to(device)
     
     # Perform inference
-    outputs = model(padded_sequences, img_x)
-    _, predicted = torch.max(outputs, 1)  
+    with torch.no_grad():  # Disable gradient calculation for inference
+        outputs = model(padded_sequences, img_x)
+        _, predicted = torch.max(outputs, 1)  
     
-    return malware_classes[predicted]
+    return malware_classes[predicted]  # Use .item() to get scalar value
 
 # Define the prediction function for Gradio
 def predict_malware(sha256_hash):
@@ -58,9 +66,9 @@ def predict_malware(sha256_hash):
         return "Hash not found in the dataset.", "", ""
 
     # Load the dataset
-    dataset = CombinedDataset(api_call_list, image_path, transforms=simple_transform ,sequence_length=config.configuration["sequence_length"])
+    dataset = CombinedDataset(api_call_list, image_path, transforms=simple_transform, sequence_length=config.configuration["sequence_length"])
     padded_sequences, img_x = next(iter(dataset))
-    img_x = img_x.unsqueeze(0) # type: ignore
+    img_x = img_x.unsqueeze(0) #type: ignore
 
     # Load the model
     model_path = "model_dump/model_malware_lstm (1).pkl"
@@ -98,14 +106,10 @@ with gr.Blocks() as demo:
             # Output for predicted malware class
             malware_output = gr.Textbox(label="Predicted Malware Class")
 
-    
-
-
     submit_button.click(
         predict_malware,
         inputs=sha256_input,
         outputs=[image_output, api_output, malware_output]
     )
 
-
 demo.launch()