harmonizing function name and signature changes

dashboards/dashboard_utils.py

@@ -53,7 +53,6 @@ def ensure_offline_available(path: Path):
 def get_available_datasets(campaign: str = "step-analysis_001") -> list[str]:
     """Scan Drive for available datasets matching pattern."""
     drive_path = Path(get_data_path()) / campaign
-    
     if not drive_path.exists():
         return []
     
@@ -65,6 +64,19 @@ def get_available_datasets(campaign: str = "step-analysis_001") -> list[str]:
             datasets.append(ds_name)
     return sorted(datasets, key=model_size_from_name)
 
+def get_available_campaigns(campaign_pattern: str = "ana-") -> list[str]:
+    """Scan Drive for available datasets matching pattern."""
+    drive_path = Path(get_data_path()) 
+    if not drive_path.exists():
+        return []
+    
+    datasets = []
+    for item in drive_path.iterdir():
+        if item.is_dir() and item.name.startswith(campaign_pattern):
+            # Extract DS_NAME by removing suffix
+            datasets.append(item.name)
+    return sorted(datasets, key=model_size_from_name)
+
 @st.cache_data
 def load_dataset_with_metadata(ds_name: str, campaign: str, hf_version: str=None):
     """Load dataset after ensuring offline availability."""

dashboards/pages/weights_dashboard.py

@@ -9,7 +9,19 @@ def weights_dashboard_app():
     plot_display = {"P(W)": "P_w", "P(λ)": "P_sv", "SVD": "SVD"}
 
     # Load data
-    df_full, metadata = load_dataset_with_metadata(ds_name='weight_study', campaign='ana-002', hf_version='ana-002')
+    available_datasets = get_available_campaigns('ana-')
+    if not available_datasets:
+        st.error(f"No datasets found.")
+        st.stop()
+
+    # Dataset dropdown
+    campaign_name = st.sidebar.selectbox(
+        "Campaign",
+        available_datasets,
+        index=0
+    )
+    
+    df_full, metadata = load_dataset_with_metadata(ds_name='weight_study', campaign=campaign_name, hf_version='ana-003')
 
     model_names = get_unique_values(df_full, "model")
     model_selected = st.sidebar.selectbox("Model", model_names)
@@ -115,6 +127,7 @@ def weights_dashboard_app():
     # Section 2: Across layers/heads
     ########################################################################
     st.header("Statistics Across Architecture")
+    df_sorted = df.sort_values(["layer", "head"])
 
     # Multi-select for statistics
     selected_stats = st.multiselect(
@@ -134,7 +147,6 @@ def weights_dashboard_app():
         )
         
         # Prepare data
-        df_sorted = df.sort_values(["layer", "head"])
         
         # Color palette
         colors = px.colors.qualitative.Plotly[:len(selected_stats)]
@@ -166,7 +178,7 @@ def weights_dashboard_app():
                 # Style the corresponding y-axis
                 axis_config = dict(
                     title=stat_display_name,
-                    titlefont=dict(color=colors[idx]),
+                    title_font=dict(color=colors[idx]),
                     tickfont=dict(color=colors[idx])
                 )
                 if use_secondary:

scripts/diff_datasets.py

@@ -56,9 +56,9 @@ def diff_datasets(name1, name2, path=None, topN=20):
     if stats1 != stats2:
         print(f"stats columns DIFFER:")
         if stats1 - stats2:
-            print(f"  Only in 1: {stats1 - stats2}")
+            print(f"  Only in reference: {stats1 - stats2}")
         if stats2 - stats1:
-            print(f"  Only in 2: {stats2 - stats1}")
+            print(f"  Only in target: {stats2 - stats1}")
         exclude_cols.update(stats1 ^ stats2)  # exclude non-common stats columns
 
     # Compare remaining metadata

scripts/merge_datasets.py

@@ -1,39 +1,76 @@
 import json
 from tqdm import tqdm
 from datasets import concatenate_datasets, load_from_disk
-from transformer_analaysis.histogram_utils import get_model_versions
+from transformer_analysis.histogram_utils import get_model_versions
 
-SUFFIX = "all_checkpoints"
 META_FILE = "metadata.json"
+META_MERGE_KEY = "merged"
 
-def merge_versions(model_name = 'pythia-70m-deduped', path = 'histos'):
+
+def write_dataset_and_metadata(ds_list, metadata, ds_name):
+    combined_ds = concatenate_datasets(ds_list)
+    combined_ds.info.description = META_FILE
+    combined_ds.save_to_disk(ds_name)
+    with open(f"{ds_name}/{META_FILE}", "w") as f:
+        json.dump(metadata, f, indent=2)
+
+def merge_versions(model_name = 'pythia-70m-deduped', path = 'histos', suffix = 'all_checkpoints'):
     ds_list = []
-    metadata_list = []
+    metadata = None
     for rev in tqdm(get_model_versions(model_name), desc=f'Processing {model_name}'):
         pattern = f"{model_name}_{rev}"
         ds = load_from_disk(f"{path}/{pattern}")
         ds_list.append(ds)
+        if metadata is None:
+            with open(f"{path}/{pattern}/{ds.info.description}") as f:
+                metadata = json.load(f)
+    write_dataset_and_metadata(ds_list, metadata,f"{path}/{model_name}_{suffix}")
+
+def merge_datasets(model_name_list, path = 'histos', out_name = 'merged', suffix=None):
+    ds_list = []
+    combined_metadata = None
+    merged_dict = {}
+    for model_name in tqdm(model_name_list, desc=f'Processing {model_name}'):
+        pattern = model_name
+        if suffix isinstance(str):
+            pattern += '_' + suffix
+        ds = load_from_disk(f"{path}/{pattern}")
+        ds_list.append(ds)
+        
+        #now the metadata
         mf = f"{path}/{pattern}/{ds.info.description}"
         with open(mf) as f:
-            x = json.load(f)
-            metadata_list.append(x)
-    return ds_list, metadata_list
-
-def write_dataset_and_metadata(ds_list, metadata_list, ds_name):
-    combined_ds = concatenate_datasets(ds_list)
-    combined_ds.info.description = META_FILE
-    combined_ds.save_to_disk(ds_name)
-    with open(f"{ds_name}/{META_FILE}", "w") as f:
-        json.dump(metadata_list[0], f, indent=2)
+            metadata = json.load(f)
+        if combined_metadata is None:
+            combined_metadata = {k: v for k, v in metadata.items() if k != merge_key}
+        if META_MERGE_KEY in metadata:  # Flatten
+            for k, v in metadata[META_MERGE_KEY].items():
+                key = k
+                while key in merged_dict:  #make key name unique
+                    key = f'{model_name}_{key}'
+                merged_dict[key] = v
+        else:
+            merged_dict[model_name] = metadata
 
+    combined_metadata.update(META_MERGE_KEY, merged_dict)
+    write_dataset_and_metadata(ds_list, combined_metadata, f"{path}/{out_name}")
 
 
 if __name__ == '__main__' :
     import argparse
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str, default="pythia-70m-deduped")
+    parser.add_argument("--model", type=str, default=None)
     parser.add_argument("--path", type=str, default='histos')
+    parser.add_argument("--out-name", type=str, default='weight_study')
+    parser.add_argument("--suffix", type=str, default='all_checkpoints')
     args = parser.parse_args()
+
+    if args.model is not None:
+        merge_versions(model_name=args.model, path=args.path, suffix=args.suffix):
     
-    ds_list, metadata_list = merge_versions(model_name=args.model, path=args.path)
-    write_dataset_and_metadata(ds_list, metadata_list, f"{args.path}/{args.model}_{SUFFIX}")
+    else:
+        import os
+        from pathlib import Path
+        path = Path(args.path)
+        model_list = [d.name for d in path.glob("*/") if args.out_name not in d.name]
+        merge_datasets(model_list, path=args.path, out_name=args.out_name)

scripts/run_model_sweep.py

@@ -1,21 +1,21 @@
 import os
-from tqdm import tqdm
-from transformers import logging as hf_logging
-from transformer_analaysisrun_weight_analysis import process_model, create_versioned_dir
-from transformer_analaysismodel_registry import MODEL_CONFIGS
+from transformer_analysis.run_weight_analysis import process_model, create_versioned_dir
+from transformer_analysis.model_registry import MODEL_CONFIGS
 
-def main(out_dir='Drive/ana-002', clobber=False):
-    cwd = os.getcwd()
+def model_sweep(model_list, out_dir='Drive/ana-002', clobber=False):
+    from tqdm import tqdm
+    from transformers import logging as hf_logging
 
     hf_logging.set_verbosity_error()
     import warnings
     warnings.filterwarnings('ignore')
 
-    #for now get mistral and llama models
-    models=[k for k in MODEL_CONFIGS.keys() if 'tral-' in k] #mistral and mixtral
-    models.extend([k for k in MODEL_CONFIGS.keys() if 'llama-' in k])
+    for model_name in tqdm(model_list):
+
+        print('\n' + '-'*40 + '\n')
+        print(f'Processing Model = {model_name}')
+        print('\n' + '-'*40 + '\n')
 
-    for model_name in tqdm(models):
         target_dir = os.path.join(out_dir, model_name)
         if os.path.exists(target_dir):
             if clobber:
@@ -24,6 +24,10 @@ def main(out_dir='Drive/ana-002', clobber=False):
                 print(f'Model = {model_name} output exists as {target_dir}. SKIPPING.')
                 continue
         process_model(model_name=model_name, cache_dir='./downloads', revision=None, out_dir=out_dir,cleanup_downloads=True)
+        print('\n' + '-'*40 + '\n')
 
 if __name__ == "__main__":
-    main()
+
+    models=[k for k in MODEL_CONFIGS.keys() if 'tral-' in k] #mistral and mixtral
+    models.extend([k for k in MODEL_CONFIGS.keys() if 'llama-' in k])
+    model_sweep(model_list=models, out_dir='Drive/ana-002', clobber=False):

src/transformer_analysis/run_weight_analysis.py

@@ -166,34 +166,22 @@ def process_model(
     logging.info(f"Performance saved to {out_dir}/logs/perf_{job_id}.json")
 
 
-def create_versioned_dir(path, name, time=False, clobber=False, increment=False):
-    """Create directory with timestamp, or numeric suffix if it exists."""
-    if time:
-        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-        base_dir = os.path.join(path, f"{name}_{timestamp}")
-    else:
-        base_dir = os.path.join(path, name)
+def create_campaign(path, name, clobber=False, logs=True):
+    base_dir = os.path.join(path, name)
+    log_dir = os.path.join(base_dir,'logs')
     
     if not os.path.exists(base_dir):
         os.makedirs(base_dir)
+        os.makedirs(log_dir)
         return base_dir
 
     elif clobber:
         shutil.rmtree(base_dir)
         os.makedirs(base_dir)
-        return base_dir
-
-    if not increment:
+        os.makedirs(log_dir)
         return base_dir
     
-    # Try numeric suffixes
-    for i in range(1, 1000):
-        versioned_dir = f"{base_dir}_{i:03d}"
-        if not os.path.exists(versioned_dir):
-            os.makedirs(versioned_dir)
-            return versioned_dir
-    
-    raise RuntimeError("Could not find available directory suffix")
+    raise RuntimeError(f"Project directory exists, rename or clobber:\n{base_dir}")
 
 
 
@@ -212,10 +200,8 @@ if __name__ == "__main__":
         print('='*20 + 'Test option selected' + '='*20)
         print('\t\t' + 'output and clobber options will be overwritten')
         args.out, args.clobber = 'test', True
-
     cwd = os.getcwd()
-    out_dir = create_versioned_dir(path=cwd, name=args.out, clobber=args.clobber)
-    log_dir = create_versioned_dir(path=out_dir, name='logs', clobber=args.clobber)
+    out_dir = create_campaign(path=cwd, name=args.out, clobber=args.clobber, logs=True)
     model_name = args.model
 
     model_config = get_model_config(args.model)