wip creates utils method to save outputs

Signed-off-by: Flavia Beo <[email protected]>

Author: flaviabeo

aiu_fms_testing_utils/utils/metrics_utils.py

@@ -1,7 +1,17 @@
 import numpy as np
 import torch
 import torch.nn as nn
+import json
 
+def save_layers_output(out_dict, file_path):
+    serializable_data = {}
+    for key, value in out_dict.items():
+        if isinstance(value, torch.Tensor):
+            serializable_data[key] = value.tolist()
+        else:
+            serializable_data[key] = value
+    with open(file_path, 'w') as f:
+        json.dump(serializable_data, f, indent=4)
 
 def abs_diff_linalg_norm(res_vector):
     """

scripts/generate_layers_metrics.py

@@ -17,7 +17,7 @@
 from aiu_fms_testing_utils.testing.validation import get_default_validation_prefix
 
 from aiu_fms_testing_utils.utils import prepare_inputs
-from aiu_fms_testing_utils.utils.metrics_utils import tensor_abs_diff, tensor_cos_sim
+from aiu_fms_testing_utils.utils.metrics_utils import tensor_abs_diff, tensor_cos_sim, save_layers_output
 
 
 logger = logging.getLogger(__name__)
@@ -390,8 +390,7 @@ def generate_layers_metrics(model_path, batch_size, seq_length, max_new_tokens):
                                             seq_length=seq_length, max_new_tokens=max_new_tokens, 
                                             tokenizer=tokenizer)
 
-    with open(os.path.join(output_path,f"{model_path}-layer-output-stack-cpu.json"), 'w') as f:
-        json.dump(layer_stack_cpu, f)
+    save_layers_output(os.path.join(output_path,f"{model_path}-layer-output-stack-cpu.json"),layer_stack_cpu)
 
     global generate_iters
     generate_iters = 0
@@ -403,8 +402,7 @@ def generate_layers_metrics(model_path, batch_size, seq_length, max_new_tokens):
                                              seq_length=seq_length, max_new_tokens=max_new_tokens, 
                                              tokenizer=tokenizer)
 
-    with open(os.path.join(output_path,f"{model_path}-layer-output-stack-gpu.json"), 'w') as f:
-        json.dump(layer_stack_cuda, f)
+    save_layers_output(os.path.join(output_path,f"{model_path}-layer-output-stack-gpu.json"), layer_stack_cuda)
 
     assert len(layer_stack_cuda.keys()) == len(layer_stack_cpu.keys())