aryan3212/clae-bengali-encoder

Continuous latent autoencoder for Bengali speech.

Architecture

• Encoder: Conformer
• Self-supervised objective: LeJEPA
• Reconstruction loss: multi-resolution STFT

Training

• Step: 170746

Config

run:
  run_id: large-2kh-packed-300k-tail-lr-1e4
  out_dir: runs
  seed: 42
  amp: true
  amp_dtype: bf16
  gpu_mem_fraction: 0.92
  wandb:
    enabled: true
    project: continuous-latent-ae
    name: large-2kh-simplified-jepa
data:
  sample_rate: 16000
  segment_seconds: 3.0
  train_manifest: staging/manifests/train.jsonl
  backend: tar
  shard_manifest: staging/packed/train/shard_manifest.json
  val_manifest: null
  num_workers: 4
  pin_memory: true
  persistent_workers: true
  prefetch_factor: 4
  shuffle_buffer_mb: 512
aug:
  waveform_aug_global:
    enabled: true
    noise_prob: 0.5
    noise_snr_min: 10.0
    noise_snr_max: 30.0
    lowpass_prob: 0.3
    lowpass_min_freq: 2000.0
    lowpass_max_freq: 8000.0
    gain_prob: 0.5
    gain_min: 0.7
    gain_max: 1.3
    clip_prob: 0.1
    clip_min: 0.5
  waveform_aug_local:
    enabled: true
    noise_prob: 0.8
    noise_snr_min: 3.0
    noise_snr_max: 15.0
    lowpass_prob: 0.6
    lowpass_min_freq: 1500.0
    lowpass_max_freq: 6000.0
    gain_prob: 0.8
    gain_min: 0.5
    gain_max: 1.7
    clip_prob: 0.3
    clip_min: 0.3
  waveform_aug_local_mask:
    enabled: true
    ratio: 0.5
    min_span_frames: 4
    max_span_frames: 24
  frontend_frame_local_mask:
    enabled: true
    ratio: 0.5
    min_span_frames: 4
    max_span_frames: 24
  frontend_frame_noise:
    enabled: false
    std: 0.0
  decoder_input_mask:
    enabled: true
    ratio: 0.5
    min_span_frames: 4
    max_span_frames: 24
  decoder_input_noise:
    enabled: true
    std: 0.03
model:
  frontend:
    channels:
    - 128
    - 256
    - 384
    - 512
    - 512
    kernels:
    - 10
    - 8
    - 8
    - 4
    - 4
    strides:
    - 5
    - 4
    - 4
    - 4
    - 4
    groups: 4
  encoder:
    encoder_type: fastconformer
    d_model: 256
    n_layers: 8
    num_heads: 8
    feedforward_dim: 1024
    dropout: 0.1
    cnn_module_kernel: 9
    use_se: true
    xscaling: false
    mhc:
      enabled: true
      num_streams: 2
      start_layer: 2
      period: 3
      sinkhorn_iters: 10
      tau: 0.05
      dropout: 0.0
      identity_mix: true
      alpha_init: 0.01
  decoder:
    channels: 768
    up_strides:
    - 4
    - 4
    - 4
    - 4
    - 5
    up_kernels:
    - 8
    - 8
    - 8
    - 8
    - 10
    res_blocks_per_up: 2
    res_dilations:
    - 1
    - 3
    - 9
    film_hidden: 192
  projector:
    hidden_dim: 512
    output_dim: 64
    n_hidden_layers: 1
loss:
  recon_type: mel
  recon_views: global
  recon_weight: 1.0
  recon_log_start_step: 10
  reg_type: visreg
  stft:
    fft_sizes:
    - 256
    - 512
    - 1024
    - 2048
    hop_ratio: 0.25
    win_ratio: 1.0
    center: true
    window: hann
    logmag_eps: 0.001
    sc_weight: 0.01
    mag_weight: 10.0
    logmag_weight: 10.0
  mel:
    n_mels: 80
    n_fft: 1024
    hop_length: 256
    win_length: 1024
    fmin: 0.0
    fmax: null
    window: hann
    logmag_eps: 0.001
    sc_weight: 0.0
    mag_weight: 1.0
    logmag_weight: 1.0
  jepa:
    weight: 0.3
    num_globals: 2
    num_locals: 4
  sigreg:
    weight: 0.05
    num_slices: 1024
    t_max: 5.0
    n_points: 17
  visreg:
    weight: 0.7
    num_projections: 256
  adv:
    enabled: false
    adv_weight: 0.3
    fm_weight: 0.5
    adv_start_step: 65000
    fm_start_step: 65000
    lr: 0.0001
    betas:
    - 0.8
    - 0.99
    periods:
    - 2
    - 3
    - 5
    - 7
    - 11
    disc_channels:
    - 32
    - 64
    - 128
    - 192
    loss_type: lsgan
    adaptive: false
    adaptive_max: 10.0
optim:
  lr: 0.0001
  betas:
  - 0.9
  - 0.999
  eps: 1.0e-08
  weight_decay: 1.0e-05
  scheduler:
    warmup_steps: 5000
    total_steps: 300000
    min_lr_ratio: 0.0
  grad_clip: 1.0
train:
  batch_size: 42
  grad_accum_steps: 4
  max_steps: 300000
  log_interval_steps: 50
  eval_interval_steps: 5000
  save_interval_steps: 25000
  probe_interval_steps: 5000
  val_batches: 10
eval:
  enabled: true
  emotion:
    enabled: false
    train_manifest: null
    dev_manifest: null
    label_key: emotion
    steps: 2000
    batch_size: 64
    segment_seconds: null
  gender:
    enabled: false
    train_manifest: null
    dev_manifest: null
    label_key: gender
    steps: 1500
    batch_size: 64
    segment_seconds: null
  asr:
    enabled: true
    train_manifest: null
    dev_manifest: null
    text_key: text
    steps: 1000
    batch_size: 16
    segment_seconds: 15.0
    max_samples: 500

How to load

import torch
ckpt = torch.load('last.pt', map_location='cpu')
state_dict = ckpt['model']
cfg = ckpt['cfg']