Upload utils.py

utils.py

@@ -210,7 +210,7 @@ def generate_memory_report(model, dataset, output_path: str = 'memory_report.jso
     
     Args:
         model: LMCODE model
-        dataset: Evaluation dataset
+        dataset: Evaluation dataset (list of dicts or list of tensors)
         output_path: Path to save report
     """
     report = {
@@ -220,8 +220,18 @@ def generate_memory_report(model, dataset, output_path: str = 'memory_report.jso
         'efficiency_metrics': {}
     }
     
-    # Analyze memory
-    test_sequences = [d['input_ids'] for d in dataset[:10]]  # Sample
+    # Analyze memory - handle different dataset formats
+    if isinstance(dataset, list):
+        if len(dataset) > 0:
+            if isinstance(dataset[0], dict):
+                test_sequences = [d['input_ids'] for d in dataset[:10]]
+            else:
+                test_sequences = dataset[:10]
+        else:
+            test_sequences = []
+    else:
+        test_sequences = []
+    
     memory_analysis = analyze_memory_capacity(model, test_sequences)
     report['memory_analysis'] = memory_analysis
     
@@ -411,4 +421,4 @@ class MemoryMonitor:
         if save_path:
             plt.savefig(save_path, dpi=150, bbox_inches='tight')
         
-        return fig
+        return fig