Update app.py

app.py

@@ -403,19 +403,16 @@ def process_hdf5_file(uploaded_file, percentage):
         # Step 4: Load the model from lwm_model module
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Loading the LWM model on {device}...")
-        model = lwm_model.LWM.from_pretrained(device=device).float()
-        #for name, param in model.state_dict().items():
-        #    print(f"Layer: {name} | Weights: {param}")
+        model = lwm_model.lwm.from_pretrained(device=device).float()
 
         # Step 5: Load the HDF5 file and extract the channels and labels
         with h5py.File(uploaded_file.name, 'r') as f:
-            channels = np.array(f['channels']).astype(np.complex64) #astype(np.float32)  # Assuming 'channels' dataset in the HDF5 file
-            labels = np.array(f['labels']).astype(np.int32)  # Assuming 'labels' dataset in the HDF5 file
+            channels = np.array(f['channels']).astype(np.complex64)
+            labels = np.array(f['labels']).astype(np.int32)  
         print(f"Loaded dataset with {channels.shape[0]} samples.")
 
         # Step 7: Tokenize the data using the tokenizer from input_preprocess
         preprocessed_chs = input_preprocess.tokenizer(manual_data=channels)
-        #print(preprocessed_chs[0][0][-1]) #CORRECT
 
         # Step 7: Perform inference using the functions from inference.py
         output_emb = inference.lwm_inference(preprocessed_chs, 'cls_emb', model, device)
@@ -424,8 +421,6 @@ def process_hdf5_file(uploaded_file, percentage):
         print(f"Output Embeddings Shape: {output_emb.shape}")
         print(f"Output Raw Shape: {output_raw.shape}")
 
-        #print(f'percentage_idx: {percentage_idx}')
-        #print(f'percentage_value: {percentage_values_los[percentage_idx]}')
         print(f'percentage_value: {percentage}')
         train_data_emb, test_data_emb, train_data_raw, test_data_raw, train_labels, test_labels = identical_train_test_split(output_emb.view(len(output_emb),-1),
                                                                                                                              output_raw.view(len(output_raw),-1),
@@ -438,8 +433,7 @@ def process_hdf5_file(uploaded_file, percentage):
         print(f'test_data_emb: {test_data_emb.shape}')
         pred_raw = classify_based_on_distance(train_data_raw, train_labels, test_data_raw)
         pred_emb = classify_based_on_distance(train_data_emb, train_labels, test_data_emb)
-        #print(f'pred_emb: {pred_emb}')
-        #print(f'actual labels: {test_labels}')
+        
         # Step 9: Generate confusion matrices for both raw and embeddings
         raw_cm_image = plot_confusion_matrix(test_labels, pred_raw, title="Confusion Matrix (Raw Channels)")
         emb_cm_image = plot_confusion_matrix(test_labels, pred_emb, title="Confusion Matrix (Embeddings)")