Zero-shot post fixes new

Author: gilkeren1

src/fairseq2/recipes/wav2vec2/asr/_train.py

@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import socket
-
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import cast, final, Literal
@@ -173,6 +172,12 @@ class Wav2Vec2AsrTrainDatasetSection(DatasetSection):
     max_num_batches: int | None = None
     """The maximum number of batches for the dataloader to return."""
 
+    always_read_tsv: bool = False
+    """If ``True``, always read the TSV manifest, regardless of whether parquet datasets exist."""
+
+    extras: dict[str, object] = field(default_factory=dict)
+    """The dataset-specific extra options."""
+
     # Upsampling
     beta_corpus: float | None = None
     beta_language: float | None = None
@@ -186,14 +191,13 @@ class Wav2Vec2AsrTrainDatasetSection(DatasetSection):
     spec_aug_time_mask_param: int = 80
     """Maximum time mask length."""
 
-    always_read_tsv: bool = False
-    """If ``True``, always read the TSV manifest, regardless of whether parquet datasets exist."""
-
-    extras: dict[str, object] = field(default_factory=dict)
-    """The dataset-specific extra options."""
-
+    # Zero/Few-shot
     n_context_examples: int = 0
     """The number of context examples to use when providing context."""
+    bucket_size_train: int = 2000
+    """Minimum size of pool for choosing context examples, for training set."""
+    bucket_size_eval: int = 30
+    """Minimum size of pool for choosing context examples, for eval sets."""
 
 
 @dataclass(kw_only=True)
@@ -435,7 +439,7 @@ def load_wav2vec2_asr_trainer(
         spec_aug_time_mask_param=config.dataset.spec_aug_time_mask_param,
         max_num_batches=config.dataset.max_num_batches,
         n_context_examples=config.dataset.n_context_examples,
-        bucket_size=2000,
+        bucket_size=config.dataset.bucket_size_train,
         deterministic_context=False,
     )
 
@@ -482,7 +486,7 @@ def load_wav2vec2_asr_trainer(
             max_num_batches=config.dataset.max_num_batches,
             extras=config.dataset.extras,
             n_context_examples=config.dataset.n_context_examples,
-            bucket_size=30,
+            bucket_size=config.dataset.bucket_size_eval,
             deterministic_context=True,
         )