Upload folder using huggingface_hub

L0-10/config.yaml

@@ -0,0 +1,201 @@
+task_name: train
+tags:
+- music4all
+- muq
+- batchtopk
+- sae
+train: true
+test: false
+ckpt_path: null
+seed: 42
+resume: false
+optimized_metrics: null
+data:
+  name: sae
+  _target_: steerable_retrieval.dataloading.dataloaders.TextAudioDataModule
+  preextracted_features: true
+  target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+  target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000} *
+    10}
+  return_audio: true
+  return_text: true
+  datasets:
+    music4all:
+      _target_: steerable_retrieval.dataloading.datasets.TextAudioDataset
+      get_annotations_function: steerable_retrieval.dataloading.dataloaders.get_music4all_annotations
+      task_kwargs:
+        manifest_csv: /gpfs/scratch/acw749/datasets/music4all_instruct/music4all_v1/embeddings/embedding_lookup_manifest.csv
+      split: keep
+      return_audio: true
+      return_text: true
+      preextracted_features: ${data.preextracted_features}
+      target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+      target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000}
+        * 10}
+      root_dir: null
+      new_dir: null
+  batch_size: 512
+  num_workers: 8
+model:
+  sae_encoder:
+    _target_: steerable_retrieval.models.sae.encoders.BatchTopKSAEEncoder
+    act_size: 512
+    dict_size: 4096
+    top_k: 10
+    threshold_momentum: 0.9
+    input_unit_norm: true
+  sae_decoder:
+    _target_: steerable_retrieval.models.sae.decoders.SAEDecoder
+    dict_size: 4096
+    act_size: 512
+    input_unit_norm: true
+    output_unit_norm: true
+  loss_fn:
+    _target_: steerable_retrieval.models.utils.losses.TopKSAELoss
+    l1_coeff: 0
+    l2_coeff: 1.0
+    aux_penalty: 0.05
+    top_k_aux: 512
+    n_batches_to_dead: 128
+    orthogonal_loss: 0
+  audio_encoder:
+    _target_: torch.nn.Identity
+  text_encoder:
+    _target_: steerable_retrieval.models.encoders.muq.MuQTextEncoder
+    model_name: OpenMuQ/MuQ-MuLan-large
+    device: cuda
+    freeze: true
+    ckpt_path: null
+  _target_: steerable_retrieval.models.sae.sae.LightningSAE
+  ckpt_path: null
+  freeze: false
+  optimizer:
+    _target_: torch.optim.AdamW
+    _partial_: true
+    lr: 0.0001
+    weight_decay: 0.0001
+    betas:
+    - 0.9
+    - 0.999
+  scheduler:
+    _target_: steerable_retrieval.models.utils.schedulers.CosineDecayWithLinearWarmup
+    _partial_: true
+    warmup_steps: 10000
+    max_steps: ${trainer.max_steps}
+    base_lr: 0.0001
+    final_lr: 1.0e-05
+  preextracted_features: true
+  data_mixing_ratio: 0.0
+callbacks:
+  model_summary:
+    _target_: lightning.pytorch.callbacks.RichModelSummary
+    max_depth: 1
+  lr_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+  model_checkpoint:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    dirpath: ${paths.ckpt_dir}
+    verbose: true
+    save_last: true
+    save_top_k: 1
+    mode: max
+    auto_insert_metric_name: true
+    save_weights_only: false
+    every_n_epochs: 1
+  progress_bar:
+    _target_: lightning.pytorch.callbacks.RichProgressBar
+  energy:
+    _target_: steerable_retrieval.callbacks.energy.EnergyCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 512
+    prefix: energy
+  modality:
+    _target_: steerable_retrieval.callbacks.energy.ModalityScoreCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    eps: 1.0e-12
+    prefix: modality
+  modality_classifier:
+    _target_: steerable_retrieval.callbacks.energy.ModalityClassifierCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    prefix: modality_classifier
+    max_samples: 10000
+    test_fraction: 0.2
+  cknna:
+    _target_: steerable_retrieval.callbacks.alignment.CKNNACallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 10
+    max_samples: 2048
+    distance_agnostic: false
+    unbiased: true
+    prefix: CKNNA
+logger:
+  wandb:
+    _target_: lightning.pytorch.loggers.wandb.WandbLogger
+    save_dir: ${paths.output_dir}/wandb/
+    offline: false
+    project: steerable-retrieval
+    log_model: false
+    save_code: true
+    prefix: ''
+    entity: jul-guinot
+    tags: ${tags}
+    job_type: ''
+trainer:
+  _target_: lightning.pytorch.trainer.Trainer
+  default_root_dir: ${paths.output_dir}
+  min_epochs: 1
+  max_steps: 500000
+  accelerator: gpu
+  devices: 1
+  strategy: auto
+  precision: 16-mixed
+  check_val_every_n_epoch: 1
+  deterministic: false
+  gradient_clip_val: 1.0
+  accumulate_grad_batches: 4
+  limit_val_batches: 100
+paths:
+  root_dir: .
+  data_dir: ${paths.root_dir}/data
+  log_dir: ${paths.root_dir}/logs/
+  pretrained: ${paths.root_dir}/pretrained
+  output_dir: ${dora:xp.folder}
+  ckpt_dir: ${paths.output_dir}/checkpoints
+extras:
+  ignore_warnings: false
+  enforce_tags: false
+  print_config: true
+dora:
+  dir: ${paths.log_dir}
+  exclude:
+  - callbacks.*
+  - data.local_dir
+  - data.num_workers
+  - data.dataloader_kwargs.pin_memory
+  - data.dataloader_kwargs.prefetch_factor
+  - data.dataloader.persistent_workers
+  - extras.*
+  - logger.*
+  - model.compile
+  - paths.*
+  - trainer.accelerator
+  - trainer.devices
+  - trainer.check_val_every_n_epoch
+  - trainer.default_root_dir
+  - trainer.min_epochs
+  - trainer.num_sanity_val_steps
+  - trainer.log_every_n_steps
+  - resume
+  git_save: false

L0-10/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa1c550be2a87b8e95bfb33fdd14be0d6454ee7db11c2afbf0265f4395e44696
+size 50364201

L0-100/config.yaml

@@ -0,0 +1,201 @@
+task_name: train
+tags:
+- music4all
+- muq
+- batchtopk
+- sae
+train: true
+test: false
+ckpt_path: null
+seed: 42
+resume: false
+optimized_metrics: null
+data:
+  name: sae
+  _target_: steerable_retrieval.dataloading.dataloaders.TextAudioDataModule
+  preextracted_features: true
+  target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+  target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000} *
+    10}
+  return_audio: true
+  return_text: true
+  datasets:
+    music4all:
+      _target_: steerable_retrieval.dataloading.datasets.TextAudioDataset
+      get_annotations_function: steerable_retrieval.dataloading.dataloaders.get_music4all_annotations
+      task_kwargs:
+        manifest_csv: /gpfs/scratch/acw749/datasets/music4all_instruct/music4all_v1/embeddings/embedding_lookup_manifest.csv
+      split: keep
+      return_audio: true
+      return_text: true
+      preextracted_features: ${data.preextracted_features}
+      target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+      target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000}
+        * 10}
+      root_dir: null
+      new_dir: null
+  batch_size: 512
+  num_workers: 8
+model:
+  sae_encoder:
+    _target_: steerable_retrieval.models.sae.encoders.BatchTopKSAEEncoder
+    act_size: 512
+    dict_size: 4096
+    top_k: 100
+    threshold_momentum: 0.9
+    input_unit_norm: true
+  sae_decoder:
+    _target_: steerable_retrieval.models.sae.decoders.SAEDecoder
+    dict_size: 4096
+    act_size: 512
+    input_unit_norm: true
+    output_unit_norm: true
+  loss_fn:
+    _target_: steerable_retrieval.models.utils.losses.TopKSAELoss
+    l1_coeff: 0
+    l2_coeff: 1.0
+    aux_penalty: 0.05
+    top_k_aux: 512
+    n_batches_to_dead: 128
+    orthogonal_loss: 0
+  audio_encoder:
+    _target_: torch.nn.Identity
+  text_encoder:
+    _target_: steerable_retrieval.models.encoders.muq.MuQTextEncoder
+    model_name: OpenMuQ/MuQ-MuLan-large
+    device: cuda
+    freeze: true
+    ckpt_path: null
+  _target_: steerable_retrieval.models.sae.sae.LightningSAE
+  ckpt_path: null
+  freeze: false
+  optimizer:
+    _target_: torch.optim.AdamW
+    _partial_: true
+    lr: 0.0001
+    weight_decay: 0.0001
+    betas:
+    - 0.9
+    - 0.999
+  scheduler:
+    _target_: steerable_retrieval.models.utils.schedulers.CosineDecayWithLinearWarmup
+    _partial_: true
+    warmup_steps: 10000
+    max_steps: ${trainer.max_steps}
+    base_lr: 0.0001
+    final_lr: 1.0e-05
+  preextracted_features: true
+  data_mixing_ratio: 0.0
+callbacks:
+  model_summary:
+    _target_: lightning.pytorch.callbacks.RichModelSummary
+    max_depth: 1
+  lr_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+  model_checkpoint:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    dirpath: ${paths.ckpt_dir}
+    verbose: true
+    save_last: true
+    save_top_k: 1
+    mode: max
+    auto_insert_metric_name: true
+    save_weights_only: false
+    every_n_epochs: 1
+  progress_bar:
+    _target_: lightning.pytorch.callbacks.RichProgressBar
+  energy:
+    _target_: steerable_retrieval.callbacks.energy.EnergyCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 512
+    prefix: energy
+  modality:
+    _target_: steerable_retrieval.callbacks.energy.ModalityScoreCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    eps: 1.0e-12
+    prefix: modality
+  modality_classifier:
+    _target_: steerable_retrieval.callbacks.energy.ModalityClassifierCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    prefix: modality_classifier
+    max_samples: 10000
+    test_fraction: 0.2
+  cknna:
+    _target_: steerable_retrieval.callbacks.alignment.CKNNACallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 10
+    max_samples: 2048
+    distance_agnostic: false
+    unbiased: true
+    prefix: CKNNA
+logger:
+  wandb:
+    _target_: lightning.pytorch.loggers.wandb.WandbLogger
+    save_dir: ${paths.output_dir}/wandb/
+    offline: false
+    project: steerable-retrieval
+    log_model: false
+    save_code: true
+    prefix: ''
+    entity: jul-guinot
+    tags: ${tags}
+    job_type: ''
+trainer:
+  _target_: lightning.pytorch.trainer.Trainer
+  default_root_dir: ${paths.output_dir}
+  min_epochs: 1
+  max_steps: 500000
+  accelerator: gpu
+  devices: 1
+  strategy: auto
+  precision: 16-mixed
+  check_val_every_n_epoch: 1
+  deterministic: false
+  gradient_clip_val: 1.0
+  accumulate_grad_batches: 4
+  limit_val_batches: 100
+paths:
+  root_dir: .
+  data_dir: ${paths.root_dir}/data
+  log_dir: ${paths.root_dir}/logs/
+  pretrained: ${paths.root_dir}/pretrained
+  output_dir: ${dora:xp.folder}
+  ckpt_dir: ${paths.output_dir}/checkpoints
+extras:
+  ignore_warnings: false
+  enforce_tags: false
+  print_config: true
+dora:
+  dir: ${paths.log_dir}
+  exclude:
+  - callbacks.*
+  - data.local_dir
+  - data.num_workers
+  - data.dataloader_kwargs.pin_memory
+  - data.dataloader_kwargs.prefetch_factor
+  - data.dataloader.persistent_workers
+  - extras.*
+  - logger.*
+  - model.compile
+  - paths.*
+  - trainer.accelerator
+  - trainer.devices
+  - trainer.check_val_every_n_epoch
+  - trainer.default_root_dir
+  - trainer.min_epochs
+  - trainer.num_sanity_val_steps
+  - trainer.log_every_n_steps
+  - resume
+  git_save: false

L0-100/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e3b758495bd98e5225e42efc92efb411b54ffd8f7525eeba168f5e874b93128
+size 50364201

L0-20/config.yaml

@@ -0,0 +1,201 @@
+task_name: train
+tags:
+- music4all
+- muq
+- batchtopk
+- sae
+train: true
+test: false
+ckpt_path: null
+seed: 42
+resume: false
+optimized_metrics: null
+data:
+  name: sae
+  _target_: steerable_retrieval.dataloading.dataloaders.TextAudioDataModule
+  preextracted_features: true
+  target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+  target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000} *
+    10}
+  return_audio: true
+  return_text: true
+  datasets:
+    music4all:
+      _target_: steerable_retrieval.dataloading.datasets.TextAudioDataset
+      get_annotations_function: steerable_retrieval.dataloading.dataloaders.get_music4all_annotations
+      task_kwargs:
+        manifest_csv: /gpfs/scratch/acw749/datasets/music4all_instruct/music4all_v1/embeddings/embedding_lookup_manifest.csv
+      split: keep
+      return_audio: true
+      return_text: true
+      preextracted_features: ${data.preextracted_features}
+      target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+      target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000}
+        * 10}
+      root_dir: null
+      new_dir: null
+  batch_size: 512
+  num_workers: 8
+model:
+  sae_encoder:
+    _target_: steerable_retrieval.models.sae.encoders.BatchTopKSAEEncoder
+    act_size: 512
+    dict_size: 4096
+    top_k: 20
+    threshold_momentum: 0.9
+    input_unit_norm: true
+  sae_decoder:
+    _target_: steerable_retrieval.models.sae.decoders.SAEDecoder
+    dict_size: 4096
+    act_size: 512
+    input_unit_norm: true
+    output_unit_norm: true
+  loss_fn:
+    _target_: steerable_retrieval.models.utils.losses.TopKSAELoss
+    l1_coeff: 0
+    l2_coeff: 1.0
+    aux_penalty: 0.05
+    top_k_aux: 512
+    n_batches_to_dead: 128
+    orthogonal_loss: 0
+  audio_encoder:
+    _target_: torch.nn.Identity
+  text_encoder:
+    _target_: steerable_retrieval.models.encoders.muq.MuQTextEncoder
+    model_name: OpenMuQ/MuQ-MuLan-large
+    device: cuda
+    freeze: true
+    ckpt_path: null
+  _target_: steerable_retrieval.models.sae.sae.LightningSAE
+  ckpt_path: null
+  freeze: false
+  optimizer:
+    _target_: torch.optim.AdamW
+    _partial_: true
+    lr: 0.0001
+    weight_decay: 0.0001
+    betas:
+    - 0.9
+    - 0.999
+  scheduler:
+    _target_: steerable_retrieval.models.utils.schedulers.CosineDecayWithLinearWarmup
+    _partial_: true
+    warmup_steps: 10000
+    max_steps: ${trainer.max_steps}
+    base_lr: 0.0001
+    final_lr: 1.0e-05
+  preextracted_features: true
+  data_mixing_ratio: 0.0
+callbacks:
+  model_summary:
+    _target_: lightning.pytorch.callbacks.RichModelSummary
+    max_depth: 1
+  lr_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+  model_checkpoint:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    dirpath: ${paths.ckpt_dir}
+    verbose: true
+    save_last: true
+    save_top_k: 1
+    mode: max
+    auto_insert_metric_name: true
+    save_weights_only: false
+    every_n_epochs: 1
+  progress_bar:
+    _target_: lightning.pytorch.callbacks.RichProgressBar
+  energy:
+    _target_: steerable_retrieval.callbacks.energy.EnergyCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 512
+    prefix: energy
+  modality:
+    _target_: steerable_retrieval.callbacks.energy.ModalityScoreCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    eps: 1.0e-12
+    prefix: modality
+  modality_classifier:
+    _target_: steerable_retrieval.callbacks.energy.ModalityClassifierCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    prefix: modality_classifier
+    max_samples: 10000
+    test_fraction: 0.2
+  cknna:
+    _target_: steerable_retrieval.callbacks.alignment.CKNNACallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 10
+    max_samples: 2048
+    distance_agnostic: false
+    unbiased: true
+    prefix: CKNNA
+logger:
+  wandb:
+    _target_: lightning.pytorch.loggers.wandb.WandbLogger
+    save_dir: ${paths.output_dir}/wandb/
+    offline: false
+    project: steerable-retrieval
+    log_model: false
+    save_code: true
+    prefix: ''
+    entity: jul-guinot
+    tags: ${tags}
+    job_type: ''
+trainer:
+  _target_: lightning.pytorch.trainer.Trainer
+  default_root_dir: ${paths.output_dir}
+  min_epochs: 1
+  max_steps: 500000
+  accelerator: gpu
+  devices: 1
+  strategy: auto
+  precision: 16-mixed
+  check_val_every_n_epoch: 1
+  deterministic: false
+  gradient_clip_val: 1.0
+  accumulate_grad_batches: 4
+  limit_val_batches: 100
+paths:
+  root_dir: .
+  data_dir: ${paths.root_dir}/data
+  log_dir: ${paths.root_dir}/logs/
+  pretrained: ${paths.root_dir}/pretrained
+  output_dir: ${dora:xp.folder}
+  ckpt_dir: ${paths.output_dir}/checkpoints
+extras:
+  ignore_warnings: false
+  enforce_tags: false
+  print_config: true
+dora:
+  dir: ${paths.log_dir}
+  exclude:
+  - callbacks.*
+  - data.local_dir
+  - data.num_workers
+  - data.dataloader_kwargs.pin_memory
+  - data.dataloader_kwargs.prefetch_factor
+  - data.dataloader.persistent_workers
+  - extras.*
+  - logger.*
+  - model.compile
+  - paths.*
+  - trainer.accelerator
+  - trainer.devices
+  - trainer.check_val_every_n_epoch
+  - trainer.default_root_dir
+  - trainer.min_epochs
+  - trainer.num_sanity_val_steps
+  - trainer.log_every_n_steps
+  - resume
+  git_save: false

L0-20/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffe764c7c57c586a4630a45b5e5bc6deb4eeb8df12af217dc3c2ffcadcacb906
+size 50364201

L0-5/config.yaml

@@ -0,0 +1,201 @@
+task_name: train
+tags:
+- music4all
+- muq
+- batchtopk
+- sae
+train: true
+test: false
+ckpt_path: null
+seed: 42
+resume: false
+optimized_metrics: null
+data:
+  name: sae
+  _target_: steerable_retrieval.dataloading.dataloaders.TextAudioDataModule
+  preextracted_features: true
+  target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+  target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000} *
+    10}
+  return_audio: true
+  return_text: true
+  datasets:
+    music4all:
+      _target_: steerable_retrieval.dataloading.datasets.TextAudioDataset
+      get_annotations_function: steerable_retrieval.dataloading.dataloaders.get_music4all_annotations
+      task_kwargs:
+        manifest_csv: /gpfs/scratch/acw749/datasets/music4all_instruct/music4all_v1/embeddings/embedding_lookup_manifest.csv
+      split: keep
+      return_audio: true
+      return_text: true
+      preextracted_features: ${data.preextracted_features}
+      target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+      target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000}
+        * 10}
+      root_dir: null
+      new_dir: null
+  batch_size: 512
+  num_workers: 8
+model:
+  sae_encoder:
+    _target_: steerable_retrieval.models.sae.encoders.BatchTopKSAEEncoder
+    act_size: 512
+    dict_size: 4096
+    top_k: 5
+    threshold_momentum: 0.9
+    input_unit_norm: true
+  sae_decoder:
+    _target_: steerable_retrieval.models.sae.decoders.SAEDecoder
+    dict_size: 4096
+    act_size: 512
+    input_unit_norm: true
+    output_unit_norm: true
+  loss_fn:
+    _target_: steerable_retrieval.models.utils.losses.TopKSAELoss
+    l1_coeff: 0
+    l2_coeff: 1.0
+    aux_penalty: 0.05
+    top_k_aux: 512
+    n_batches_to_dead: 128
+    orthogonal_loss: 0
+  audio_encoder:
+    _target_: torch.nn.Identity
+  text_encoder:
+    _target_: steerable_retrieval.models.encoders.muq.MuQTextEncoder
+    model_name: OpenMuQ/MuQ-MuLan-large
+    device: cuda
+    freeze: true
+    ckpt_path: null
+  _target_: steerable_retrieval.models.sae.sae.LightningSAE
+  ckpt_path: null
+  freeze: false
+  optimizer:
+    _target_: torch.optim.AdamW
+    _partial_: true
+    lr: 0.0001
+    weight_decay: 0.0001
+    betas:
+    - 0.9
+    - 0.999
+  scheduler:
+    _target_: steerable_retrieval.models.utils.schedulers.CosineDecayWithLinearWarmup
+    _partial_: true
+    warmup_steps: 10000
+    max_steps: ${trainer.max_steps}
+    base_lr: 0.0001
+    final_lr: 1.0e-05
+  preextracted_features: true
+  data_mixing_ratio: 0.0
+callbacks:
+  model_summary:
+    _target_: lightning.pytorch.callbacks.RichModelSummary
+    max_depth: 1
+  lr_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+  model_checkpoint:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    dirpath: ${paths.ckpt_dir}
+    verbose: true
+    save_last: true
+    save_top_k: 1
+    mode: max
+    auto_insert_metric_name: true
+    save_weights_only: false
+    every_n_epochs: 1
+  progress_bar:
+    _target_: lightning.pytorch.callbacks.RichProgressBar
+  energy:
+    _target_: steerable_retrieval.callbacks.energy.EnergyCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 512
+    prefix: energy
+  modality:
+    _target_: steerable_retrieval.callbacks.energy.ModalityScoreCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    eps: 1.0e-12
+    prefix: modality
+  modality_classifier:
+    _target_: steerable_retrieval.callbacks.energy.ModalityClassifierCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    prefix: modality_classifier
+    max_samples: 10000
+    test_fraction: 0.2
+  cknna:
+    _target_: steerable_retrieval.callbacks.alignment.CKNNACallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 10
+    max_samples: 2048
+    distance_agnostic: false
+    unbiased: true
+    prefix: CKNNA
+logger:
+  wandb:
+    _target_: lightning.pytorch.loggers.wandb.WandbLogger
+    save_dir: ${paths.output_dir}/wandb/
+    offline: false
+    project: steerable-retrieval
+    log_model: false
+    save_code: true
+    prefix: ''
+    entity: jul-guinot
+    tags: ${tags}
+    job_type: ''
+trainer:
+  _target_: lightning.pytorch.trainer.Trainer
+  default_root_dir: ${paths.output_dir}
+  min_epochs: 1
+  max_steps: 500000
+  accelerator: gpu
+  devices: 1
+  strategy: auto
+  precision: 16-mixed
+  check_val_every_n_epoch: 1
+  deterministic: false
+  gradient_clip_val: 1.0
+  accumulate_grad_batches: 4
+  limit_val_batches: 100
+paths:
+  root_dir: .
+  data_dir: ${paths.root_dir}/data
+  log_dir: ${paths.root_dir}/logs/
+  pretrained: ${paths.root_dir}/pretrained
+  output_dir: ${dora:xp.folder}
+  ckpt_dir: ${paths.output_dir}/checkpoints
+extras:
+  ignore_warnings: false
+  enforce_tags: false
+  print_config: true
+dora:
+  dir: ${paths.log_dir}
+  exclude:
+  - callbacks.*
+  - data.local_dir
+  - data.num_workers
+  - data.dataloader_kwargs.pin_memory
+  - data.dataloader_kwargs.prefetch_factor
+  - data.dataloader.persistent_workers
+  - extras.*
+  - logger.*
+  - model.compile
+  - paths.*
+  - trainer.accelerator
+  - trainer.devices
+  - trainer.check_val_every_n_epoch
+  - trainer.default_root_dir
+  - trainer.min_epochs
+  - trainer.num_sanity_val_steps
+  - trainer.log_every_n_steps
+  - resume
+  git_save: false

L0-5/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:496f8852792593a4fec12b92d6abb420814220a752567415e7c6bd36cd1fd729
+size 50364201

L0-50/config.yaml

@@ -0,0 +1,201 @@
+task_name: train
+tags:
+- music4all
+- muq
+- batchtopk
+- sae
+train: true
+test: false
+ckpt_path: null
+seed: 42
+resume: false
+optimized_metrics: null
+data:
+  name: sae
+  _target_: steerable_retrieval.dataloading.dataloaders.TextAudioDataModule
+  preextracted_features: true
+  target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+  target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000} *
+    10}
+  return_audio: true
+  return_text: true
+  datasets:
+    music4all:
+      _target_: steerable_retrieval.dataloading.datasets.TextAudioDataset
+      get_annotations_function: steerable_retrieval.dataloading.dataloaders.get_music4all_annotations
+      task_kwargs:
+        manifest_csv: /gpfs/scratch/acw749/datasets/music4all_instruct/music4all_v1/embeddings/embedding_lookup_manifest.csv
+      split: keep
+      return_audio: true
+      return_text: true
+      preextracted_features: ${data.preextracted_features}
+      target_sr: ${oc.select:model.audio_encoder.sampling_rate,24000}
+      target_n_samples: ${eval:${oc.select:model.audio_encoder.sampling_rate,24000}
+        * 10}
+      root_dir: null
+      new_dir: null
+  batch_size: 512
+  num_workers: 8
+model:
+  sae_encoder:
+    _target_: steerable_retrieval.models.sae.encoders.BatchTopKSAEEncoder
+    act_size: 512
+    dict_size: 4096
+    top_k: 50
+    threshold_momentum: 0.9
+    input_unit_norm: true
+  sae_decoder:
+    _target_: steerable_retrieval.models.sae.decoders.SAEDecoder
+    dict_size: 4096
+    act_size: 512
+    input_unit_norm: true
+    output_unit_norm: true
+  loss_fn:
+    _target_: steerable_retrieval.models.utils.losses.TopKSAELoss
+    l1_coeff: 0
+    l2_coeff: 1.0
+    aux_penalty: 0.05
+    top_k_aux: 512
+    n_batches_to_dead: 128
+    orthogonal_loss: 0
+  audio_encoder:
+    _target_: torch.nn.Identity
+  text_encoder:
+    _target_: steerable_retrieval.models.encoders.muq.MuQTextEncoder
+    model_name: OpenMuQ/MuQ-MuLan-large
+    device: cuda
+    freeze: true
+    ckpt_path: null
+  _target_: steerable_retrieval.models.sae.sae.LightningSAE
+  ckpt_path: null
+  freeze: false
+  optimizer:
+    _target_: torch.optim.AdamW
+    _partial_: true
+    lr: 0.0001
+    weight_decay: 0.0001
+    betas:
+    - 0.9
+    - 0.999
+  scheduler:
+    _target_: steerable_retrieval.models.utils.schedulers.CosineDecayWithLinearWarmup
+    _partial_: true
+    warmup_steps: 10000
+    max_steps: ${trainer.max_steps}
+    base_lr: 0.0001
+    final_lr: 1.0e-05
+  preextracted_features: true
+  data_mixing_ratio: 0.0
+callbacks:
+  model_summary:
+    _target_: lightning.pytorch.callbacks.RichModelSummary
+    max_depth: 1
+  lr_monitor:
+    _target_: lightning.pytorch.callbacks.LearningRateMonitor
+  model_checkpoint:
+    _target_: lightning.pytorch.callbacks.ModelCheckpoint
+    dirpath: ${paths.ckpt_dir}
+    verbose: true
+    save_last: true
+    save_top_k: 1
+    mode: max
+    auto_insert_metric_name: true
+    save_weights_only: false
+    every_n_epochs: 1
+  progress_bar:
+    _target_: lightning.pytorch.callbacks.RichProgressBar
+  energy:
+    _target_: steerable_retrieval.callbacks.energy.EnergyCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 512
+    prefix: energy
+  modality:
+    _target_: steerable_retrieval.callbacks.energy.ModalityScoreCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    eps: 1.0e-12
+    prefix: modality
+  modality_classifier:
+    _target_: steerable_retrieval.callbacks.energy.ModalityClassifierCallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    prefix: modality_classifier
+    max_samples: 10000
+    test_fraction: 0.2
+  cknna:
+    _target_: steerable_retrieval.callbacks.alignment.CKNNACallback
+    enable_on_validation: true
+    enable_on_test: true
+    every_n_steps: null
+    every_n_epochs: 1
+    topk: 10
+    max_samples: 2048
+    distance_agnostic: false
+    unbiased: true
+    prefix: CKNNA
+logger:
+  wandb:
+    _target_: lightning.pytorch.loggers.wandb.WandbLogger
+    save_dir: ${paths.output_dir}/wandb/
+    offline: false
+    project: steerable-retrieval
+    log_model: false
+    save_code: true
+    prefix: ''
+    entity: jul-guinot
+    tags: ${tags}
+    job_type: ''
+trainer:
+  _target_: lightning.pytorch.trainer.Trainer
+  default_root_dir: ${paths.output_dir}
+  min_epochs: 1
+  max_steps: 500000
+  accelerator: gpu
+  devices: 1
+  strategy: auto
+  precision: 16-mixed
+  check_val_every_n_epoch: 1
+  deterministic: false
+  gradient_clip_val: 1.0
+  accumulate_grad_batches: 4
+  limit_val_batches: 100
+paths:
+  root_dir: .
+  data_dir: ${paths.root_dir}/data
+  log_dir: ${paths.root_dir}/logs/
+  pretrained: ${paths.root_dir}/pretrained
+  output_dir: ${dora:xp.folder}
+  ckpt_dir: ${paths.output_dir}/checkpoints
+extras:
+  ignore_warnings: false
+  enforce_tags: false
+  print_config: true
+dora:
+  dir: ${paths.log_dir}
+  exclude:
+  - callbacks.*
+  - data.local_dir
+  - data.num_workers
+  - data.dataloader_kwargs.pin_memory
+  - data.dataloader_kwargs.prefetch_factor
+  - data.dataloader.persistent_workers
+  - extras.*
+  - logger.*
+  - model.compile
+  - paths.*
+  - trainer.accelerator
+  - trainer.devices
+  - trainer.check_val_every_n_epoch
+  - trainer.default_root_dir
+  - trainer.min_epochs
+  - trainer.num_sanity_val_steps
+  - trainer.log_every_n_steps
+  - resume
+  git_save: false

L0-50/last.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8dce23f477346f90c2568b5adcead72b49942725c59e504f44b3bab804b012ca
+size 50364201

README.md

@@ -0,0 +1,8 @@
+---
+license: mit
+tags: [music, retrieval, sparse-autoencoder, muq]
+---
+
+# Sparse Steerable Retrieval — BatchTopK SAEs
+
+BatchTopK SAEs trained on MuQ-MuLan music4all audio embeddings (512-d), one per sparsity level L0. Each `L0-*/` holds `last.ckpt` (encoder/decoder weights) + `config.yaml`. Load via `steerable_retrieval.steer.load_steerable_sae(repo, subfolder='L0-20')`.