Update inference.py

inference.py

@@ -23,29 +23,9 @@ import numpy as np
 import warnings
 warnings.filterwarnings('ignore')
 
-def set_seed(seed=42):
-    torch.manual_seed(seed)
-    np.random.seed(seed)
-
-# Use this function at the start of your code
-set_seed(42)
-
-# Force model weights and data to float32 precision
-def cast_model_weights_to_float32(model):
-    for param in model.parameters():
-        param.data = param.data.float()  # Cast all weights to float32
-    return model
-
-# Device configuration
-device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
-if torch.cuda.is_available():
-    torch.cuda.empty_cache()
-
-def lwm_inference(preprocessed_chs, input_type, lwm_model):
+def lwm_inference(preprocessed_chs, input_type, lwm_model, device):
     
     dataset = prepare_for_LWM(preprocessed_chs, device)
-
-    lwm_model = cast_model_weights_to_float32(lwm_model)
     # Process data through LWM
     lwm_loss, embedding_data = evaluate(lwm_model, dataset)
     print(f'LWM loss: {lwm_loss:.4f}')
@@ -56,15 +36,14 @@ def lwm_inference(preprocessed_chs, input_type, lwm_model):
         embedding_data = embedding_data[:, 1:]
     
     dataset = embedding_data.float()
-    print(dataset[0][:4])    
     return dataset
 
 def prepare_for_LWM(data, device, batch_size=64, shuffle=False):
 
     input_ids, masked_tokens, masked_pos = zip(*data)
     
-    input_ids_tensor = torch.tensor(input_ids, device=device).float()  # Explicitly cast to float32
-    masked_tokens_tensor = torch.tensor(masked_tokens, device=device).float()  # Explicitly cast to float32
+    input_ids_tensor = torch.tensor(input_ids, device=device).float() 
+    masked_tokens_tensor = torch.tensor(masked_tokens, device=device).float() 
     masked_pos_tensor = torch.tensor(masked_pos, device=device).long()
 
     dataset = TensorDataset(input_ids_tensor, masked_tokens_tensor, masked_pos_tensor)
@@ -84,16 +63,13 @@ def evaluate(model, dataloader):
             masked_tokens = batch[1]
             masked_pos = batch[2]
 
-            if idx == 0:
-                print(input_ids[0])
-                
             logits_lm, output = model(input_ids, masked_pos)
             
             output_batch_preproc = output 
             outputs.append(output_batch_preproc)
 
             loss_lm = criterionMCM(logits_lm, masked_tokens)
-            loss = loss_lm / torch.var(masked_tokens)  # Use variance for normalization
+            loss = loss_lm / torch.var(masked_tokens)  
             running_loss += loss.item()
             
     average_loss = running_loss / len(dataloader)
@@ -104,6 +80,6 @@ def evaluate(model, dataloader):
 def create_raw_dataset(data, device):
     """Create a dataset for raw channel data."""
     input_ids, _, _ = zip(*data)
-    input_data = torch.tensor(input_ids, device=device).float()[:, 1:]  # Explicitly cast to float32
+    input_data = torch.tensor(input_ids, device=device)[:, 1:]  
     return input_data.float()