Update code to pass security scan

Author: Ganyu Teng

anollm/anollm.py

@@ -22,6 +22,7 @@
 from anollm.anollm_utils import _array_to_dataframe
 from anollm.anollm_dataset import AnoLLMDataset, AnoLLMDataCollator
 
+from safetensors.torch import save_model, load_model
 
 class AnoLLM:
 	"""AnoLLM Class
@@ -265,7 +266,7 @@ def decision_function(
 				shift_attention_mask_batch = attn_mask[..., 1:].contiguous()
 
 				if feature_wise:
-					score_batch = (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).cpu().numpy() # batch * (ori_seq_len -1)
+					score_batch = (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).cpu().to(torch.float32).numpy() # batch * (ori_seq_len -1)
 
 					for i in range(len(encoded_batch)):
 						for j in range(n_col): 
@@ -274,7 +275,7 @@ def decision_function(
 							col_idx = col_indices_batch[i][j]
 							anomaly_scores[start_idx+i, col_idx, perm_idx] = score_batch[i, start_pos:end_pos].sum()
 				elif len(self.textual_columns) > 0:
-					score_batch = (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).cpu().numpy() # batch * (ori_seq_len -1)
+					score_batch = (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).cpu().to(torch.float32).numpy() # batch * (ori_seq_len -1)
 					for i in range(len(encoded_batch)):
 						score_single = 0
 						for j in range(n_col): 
@@ -287,7 +288,7 @@ def decision_function(
 								score_single += score_batch[i, start_pos:end_pos].sum()
 						anomaly_scores[start_idx+i, perm_idx] = score_single
 				else:
-					score_batch = (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).sum(1) # remove normalization
+					score_batch = (loss_fct(shift_logits.transpose(1, 2), shift_labels) * shift_attention_mask_batch).to(torch.float32).sum(1) # remove normalization
 					anomaly_scores[start_idx:end_idx, perm_idx] = score_batch.cpu().numpy()
 				start_idx = end_idx
 
@@ -309,24 +310,14 @@ def save_state_dict(self, path: str):
 		else:
 			os.mkdir(directory)
 
-		state_dict = self.model.state_dict()
-		new_state_dict = OrderedDict()
-		for k, v in state_dict.items():
-			name = k[7:]  # remove `module.`
-			new_state_dict[name] = v
-		# Save the model with the modified state dict
-		torch.save(new_state_dict, path)
+		model_to_save = self.model.module
+		save_model(model_to_save, path)
 
 	def load_from_state_dict(self, path: str):
 		"""Load AnoLLM model from state_dict
 
 		Args:
 			path: path where AnoLLM model is saved
 		"""
-		
-		if self.efficient_finetuning == 'lora':
-			self.model.to('cpu')
-			state_dict = torch.load(path, map_location=torch.device('cpu'))
-			self.model.load_state_dict(state_dict)
-		else:
-			self.model.load_state_dict(torch.load(path))
+		load_model(self.model, path)
+

requirements.txt

@@ -7,9 +7,9 @@ pandas==2.2.2
 scikit_learn==1.6.1
 scipy==1.13.1
 tqdm==4.66.4
-transformers==4.42.3
 ucimlrepo==0.0.7
 peft==0.11.1
 datasets==2.20.0
 wandb==0.17.4
-tf-keras==2.16.0
+tf-keras==2.16.0
+transformers==4.48.2

scripts/exp2-odds/run_anollm.sh

@@ -18,7 +18,6 @@ for model in 'smol' 'smol-360'; do
                                                     --batch_size $batch_size --model $model --binning standard  --wandb
         CUDA_VISIBLE_DEVICES=$INFERENCE_GPUS torchrun --nproc_per_node=$n_test_node evaluate_anollm.py --dataset $dataset --n_splits $n_splits --split_idx 0  --setting $setting\
                                                 --batch_size $eval_batch_size  --n_permutations $n_permutations --model $model --binning standard   
-        exit
         wandb offline
         for ((split_idx = 1 ; split_idx < $n_splits ; split_idx++ )); do    
             CUDA_VISIBLE_DEVICES=$TRAIN_GPUS torchrun --nproc_per_node=$n_train_node train_anollm.py --dataset $dataset --n_splits $n_splits --split_idx $split_idx  --setting $setting --max_steps 2000\

src/data_utils.py

@@ -459,7 +459,7 @@ def flatten(l):
 		for npz_file in os.listdir(dataset_root):
 			if npz_file.startswith(str(n) + '_'):
 				print(dataset_name, npz_file)
-				data = np.load(dataset_root / npz_file, allow_pickle=True)
+				data = np.load(dataset_root / npz_file, allow_pickle=False)
 				break
 		else: 
 			ValueError('{} is not found.'.format(dataset_name))
@@ -482,11 +482,11 @@ def load_adbench_data(dataset):
 		Utils().download_datasets(repo='jihulab')
 
 	if dataset == 'cardio':
-		return np.load(dataset_root / '6_cardio.npz', allow_pickle=True)
+		return np.load(dataset_root / '6_cardio.npz', allow_pickle=False)
 
 	for npz_file in os.listdir(dataset_root):
 		if dataset in npz_file.lower():
-			return np.load(dataset_root / npz_file, allow_pickle=True)
+			return np.load(dataset_root / npz_file, allow_pickle=False)
 	else: 
 		ValueError('{} is not found.'.format(dataset))