Replacing weight with multiplier

ultravox/data/dataset_config.py

@@ -12,7 +12,7 @@ class DataDictConfig(BaseModel):
     splits: List[str] = dataclasses.field(default_factory=list)
     num_samples: Optional[int] = None
     total_samples: int = 1
-    weight: float = 1.0
+    multiplier: float = 1.0
     streaming: bool = True
     user_template: str = "<|audio|>"
     assistant_template: str = "{{text}}"

ultravox/data/datasets.py

@@ -279,7 +279,7 @@ class SizedIterableDataset(abc.ABC, data.IterableDataset):
     """
 
     @abc.abstractmethod
-    def __len__(self):
+    def __len__(self) -> int:
         pass
 
 
@@ -296,16 +296,16 @@ def __init__(self, args: VoiceDatasetArgs) -> None:
         self._args = args
         self._session: Optional[requests.Session] = None
         self._rng = np.random.default_rng(self._args.shuffle_seed)
-        self._weight = 1.0  # the default weight for the dataset
+        self._multiplier = 1.0
 
     def _init_dataset(self, dataset: data.Dataset, estimated_length: int = 1) -> None:
         self._dataset = dataset
         # Only required when using epochs when training dataset.
         self._estimated_length = estimated_length
 
     @property
-    def weight(self) -> float:
-        return self._weight
+    def multiplier(self) -> float:
+        return self._multiplier
 
     def _load_audio_dataset(
         self,
@@ -373,7 +373,7 @@ def __iter__(self):
                 f"Mismatch between estimated length ({self._estimated_length}) and actual length ({actual_length}) for dataset of type {type(self._dataset)}. Make sure to update."
             )
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self._estimated_length
 
     @abc.abstractmethod
@@ -493,7 +493,7 @@ def __init__(self, estimated_length: int = 1) -> None:
     def __iter__(self):
         return iter([])
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self._estimated_length
 
 
@@ -1049,16 +1049,17 @@ def __init__(
         if self._args.shuffle:
             dataset = dataset.shuffle(seed=self._args.shuffle_seed)
 
-        if config.num_samples:
-            dataset = Range(dataset, config.num_samples, config.total_samples)
-
-        self._weight = config.weight
+        self._multiplier = config.multiplier
 
         self.user_template = config.user_template
         self.assistant_template = config.assistant_template
         self.transcript_template = config.transcript_template
 
-        super()._init_dataset(dataset, config.total_samples)
+        if config.num_samples:
+            dataset = Range(dataset, config.num_samples, config.total_samples)
+            super()._init_dataset(dataset, len(dataset))
+        else:
+            super()._init_dataset(dataset, config.total_samples)
 
     def _get_sample(self, row) -> VoiceSample:
         try:
@@ -1121,7 +1122,7 @@ class StopStrategy(str, Enum):
 
 
 class InterleaveDataset(SizedIterableDataset):
-    """Interleaves multiple IterableDataset objects based on normalized weights."""
+    """Interleaves multiple IterableDataset objects based on multiplier."""
 
     def __init__(
         self,
@@ -1142,10 +1143,12 @@ def __init__(
         self._static = static
 
         self._stop_strategy = stop_strategy
+        relative_frequencies = [
+            int(getattr(ds, "multiplier", 1.0) * float(len(ds))) for ds in datasets
+        ]
 
-        weights = [getattr(ds, "weight", 1) for ds in datasets]
-        total_weight = sum(weights)
-        self._normalized_probs = [w / total_weight for w in weights]
+        total_frequency = sum(relative_frequencies)
+        self._normalized_probs = [f / total_frequency for f in relative_frequencies]
 
     def __iter__(self):
         # If no datasets are provided, return an empty iterator
@@ -1180,9 +1183,12 @@ def __iter__(self):
                 iters[iter_index] = iter(self._datasets[iter_index])
                 yield next(iters[iter_index])
 
-    def __len__(self):
+    def __len__(self) -> int:
         # TODO: Implement the length method for different stop strategies
-        return sum(len(ds) for ds in self._datasets)
+        return sum(
+            int(getattr(ds, "multiplier", 1.0) * float(len(ds)))
+            for ds in self._datasets
+        )
 
 
 class Dataproc(SizedIterableDataset):
@@ -1198,7 +1204,7 @@ def _process(self, sample: VoiceSample) -> Dict[str, Any]:
     def __iter__(self):
         return (self._process(sample) for sample in self._dataset)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self._dataset)
 
 
@@ -1234,7 +1240,7 @@ def __iter__(self):
                 break
             yield sample
 
-    def __len__(self):
+    def __len__(self) -> int:
         return (
             self._num_samples
             if self._num_samples is not None

ultravox/data/datasets_test.py

@@ -15,14 +15,14 @@
 class FakeSizedIterableDataset(datasets.SizedIterableDataset):
     """Fake version of datasets.SizedIterableDataset"""
 
-    def __init__(self, n, start=0, weight=1, estimated_length=0):
+    def __init__(self, n, start=0, multiplier=1, estimated_length=1):
         self.data = range(start, start + n)
-        self._weight = weight
+        self._multiplier = multiplier
         self._estimated_length = estimated_length
 
     @property
-    def weight(self) -> float:
-        return self._weight
+    def multiplier(self) -> float:
+        return self._multiplier
 
     def __iter__(self):
         for sample in self.data:
@@ -135,8 +135,8 @@ def test_interleaved_never_stop():
 
 
 def test_interleaved_random():
-    ds1 = FakeSizedIterableDataset(4, weight=10)
-    ds2 = FakeSizedIterableDataset(2, start=10, weight=1)
+    ds1 = FakeSizedIterableDataset(4, multiplier=10)
+    ds2 = FakeSizedIterableDataset(2, start=10, multiplier=1)
     s = datasets.InterleaveDataset(
         [ds1, ds2],
     )

ultravox/training/configs/meta_config.yaml

@@ -1,7 +1,7 @@
 text_model: "meta-llama/Meta-Llama-3-8B-Instruct"
 audio_model: "facebook/wav2vec2-base-960h"
 
-data_sets: ["gigaspeech"]
+data_sets: [] # Can't use datasets with "multiplier" because it doesn't have a length attribute.
 val_sets: ["heysquad_human", "anyinstruct", "soda", "peoplespeech"]
 stop_strategy: "LAST_EXHAUSTED"