Skip to content

kirad.datasets_and_dataloaders

kirad.datasets_and_dataloaders

Data stage: metadata table, per-modality datasets (audio, text, cache) and UUID sampler.

KintsugiMetadata loads the CSV and labels; the audio/text datasets return one recording's windows / tokens by UUID; launcher_utils stacks them and dataset_utils.broadcast_collate batches.

Base datasets and data loaders. Currently only contains the audio loader for the Whisper Medium backbone.

INFERENCE_WINDOW_SIZE = 30 module-attribute

MAX_INFERENCE_WINDOW_OVERLAP_FRACTION = 1 / 3 module-attribute

INFERENCE_BATCH_SIZE = 16 module-attribute

DEFAULT_LABEL_COLUMN = {'depression': 'phq', 'anxiety': 'gad'} module-attribute

MERGE_ROWS_BY_SUBJECT_VALUES = Literal[False, 'many_to_one', 'many_to_many'] module-attribute

TEXT_OUTPUT_PER_SAMPLE_VALUES = Literal['cyclic', 'merge'] module-attribute

WINDOW_METHOD_VALUES = Literal['all', 'single_random', 'single_start', 'random_adjacent_pair'] module-attribute

KintsugiMetadata

Bases: BaseModel, Dataset

Metadata table indexed by UUID: loads CSV, merges KD scores, filters by duration.

Note: duration_column is required (num_windows and duration filters use it).

Class for reading and processing metadata for training.

Attributes:

Name Type Description
df_or_dataset_version DataFrame or string

If DataFrame, must contain the filename of the files being loaded in the filename column. If a string, specifies the dataset version to load from Kintsugi's standardized dataset splits.
audio_root Path or string

If a Path, defines the directory from which audio files should be loaded. If a string, defines the feature store bucket to load audio from. Available feature store bucket options are raw, resampled, resampled_and_cut_audio_with_conversational_vad_params, resampled_and_trimmed_audio_with_conversational_vad_params, resampled_and_cut_audio_with_prompt_vad_params, and resampled_and_trimmed_audio_with_prompt_vad_params.
text_root Path or string

Path to directory containing transcription json files.
data_sources string or list of strings, optional

If df_or_dataset_version defines a dataset version, the data_sources parameter defines which data sources to load. For example, if data_sources is phonic, only the phonic metadata is used in the specified Kintsugi dataset. Available data sources are phonic, sonar, sonar_general_pop, sonar_senior_pop, sonar_elderly_and_latino_pop, fda, app, and all. The sonar data source only exists in the v1 dataset and refers to the senior population data. The data from Sonar was broken into senior_general_pop and sonar_senior_pop in later datasets.
dataset_splits string or list of strings, optional

If df_or_dataset_version defines a dataset version, the dataset_split parameter defines which split to load. For example, if dataset_split is train, only the train split is loaded. Available splits are train, val, test, and all
dataset_fold (int, optional)

If df_or_dataset_version defines a dataset version, the dataset_fold parameter defines which fold to load.
label_column (dict, optional)

Specify a dictionary with (task name, label_column) key/value pairs for the column containing the ground truth scores for each task
duration_column (string, optional)

Specify column containing the audio duration.
merge_rows_by_subject Literal[False, 'many_to_one', 'many_to_many']

If False, leave data as is, without combining rows corresponding to the same subject in any way -- this should always be the choice for val and test.

Otherwise, replace all the rows corresponding to a subject with either a single row ("many_to_one") or the same number of rows ("many_to_many"), as follows. In both cases the audio is concatenated across all these rows before windows / features are extracted. In the "many_to_one" case the text output is the text of a random one of the subjects' rows, chosen independently at each epoch. In the "many_to_many" case the text isn't changed -- if a subject had three streams before merging with texts A, B, and C, the dataset will still contain three streams, one with each of texts A, B, and C. Only the audio is concatenated, so it no longer corresponds directly to the text.

The way these options interact with window_method to affect text output is slightly counterintuitive, but is chosen to make score variance loss work effectively. In case merge_rows_by_subject is not False and window_method is set to return multiple windows, then the text for the first window is chosen as described above and the text for subsequent windows cycle through other text from that subject. Continuing the above three-stream example, if the text output for the first window would have been B, the text output for the second window would be C, for the third window it would be A, for the fourth B, etc.
window_duration int

The number of seconds to grab from each audio file before feature extraction.
max_overlap_frac (float, optional)

Specify how much the windows overlap as a fraction of the window length. Default: 0.0.
score_targets_csv Optional[str]

If present, filename of a CSV file with columns "uuid" and f"scores_{task}" for each task in label_type. These scores will be packaged with the quantized labels for the elements of this dataset, so they can be used by the model e.g. as targets for knowledge distillation. This CSV is typically generated by: (1) Running evaluation of a model using eval_set = "train". (2) Identifying the WANDB_RUN_ID of this eval job (3) Running kirad.utils.get_dataframe_with_scores(WANDB_RUN_ID, "test").to_csv(OUT_CSV) (note this must be called with split="test", not split="train" -- see the comment in get_dataframe_with_scores for a description of this inconsistency.) but could alternatively be generated by some other ad hoc process, e.g. combining scores of multiple models in some way.
Source code in kirad/datasets_and_dataloaders.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
class KintsugiMetadata(BaseModel, Dataset):
    """Metadata table indexed by UUID: loads CSV, merges KD scores, filters by duration.

    Note: duration_column is required (num_windows and duration filters use it).

    Class for reading and processing metadata for training.

    Attributes
    ----------
    df_or_dataset_version : pd.DataFrame or string
        If DataFrame, must contain the filename of the files being loaded in
        the `filename` column. If a string, specifies the dataset version to
        load from Kintsugi's standardized dataset splits.
    audio_root : Path or string
        If a Path, defines the directory from which audio files should be
        loaded. If a string, defines the feature store bucket to load audio
        from. Available feature store bucket options are `raw`, `resampled`,
        `resampled_and_cut_audio_with_conversational_vad_params`,
        `resampled_and_trimmed_audio_with_conversational_vad_params`,
        `resampled_and_cut_audio_with_prompt_vad_params`, and
        `resampled_and_trimmed_audio_with_prompt_vad_params`.
    text_root : Path or string
        Path to directory containing transcription json files.
    data_sources : string or list of strings, optional
        If `df_or_dataset_version` defines a dataset version, the
        `data_sources` parameter defines which data sources to load. For
        example, if `data_sources` is `phonic`, only the `phonic` metadata
        is used in the specified Kintsugi dataset. Available data sources are
        `phonic`, `sonar`, `sonar_general_pop`, `sonar_senior_pop`,
        `sonar_elderly_and_latino_pop`, `fda`, `app`, and `all`. The `sonar` data
        source only exists in the v1 dataset and refers to the senior population data.
        The data from Sonar was broken into `senior_general_pop` and
        `sonar_senior_pop` in later datasets.
    dataset_splits : string or list of strings, optional
        If `df_or_dataset_version` defines a dataset version, the
        `dataset_split` parameter defines which split to load. For example,
        if `dataset_split` is `train`, only the train split is loaded. Available
        splits are `train`, `val`, `test`, and `all`
    dataset_fold : int, optional
        If `df_or_dataset_version` defines a dataset version, the
        `dataset_fold` parameter defines which fold to load.
    label_column : dict, optional
        Specify a dictionary with (task name, label_column) key/value pairs for the
        column containing the ground truth scores for each task
    duration_column : string, optional
        Specify column containing the audio duration.
    merge_rows_by_subject : Literal[False, "many_to_one", "many_to_many"]
        If False, leave data as is, without combining rows corresponding to the same subject in any
        way -- this should always be the choice for val and test.

        Otherwise, replace all the rows corresponding to a subject with either a single row
        ("many_to_one") or the same number of rows ("many_to_many"), as follows. In both cases the
        audio is concatenated across all these rows before windows / features are extracted. In the
        "many_to_one" case the text output is the text of a random one of the subjects' rows,
        chosen independently at each epoch. In the "many_to_many" case the text isn't changed -- if
        a subject had three streams before merging with texts A, B, and C, the dataset will still
        contain three streams, one with each of texts A, B, and C. Only the audio is concatenated,
        so it no longer corresponds directly to the text.

        The way these options interact with `window_method` to affect text output is slightly
        counterintuitive, but is chosen to make score variance loss work effectively. In case
        `merge_rows_by_subject` is not False and `window_method` is set to return multiple windows,
        then the text for the first window is chosen as described above and the text for subsequent
        windows cycle through other text from that subject. Continuing the above three-stream
        example, if the text output for the first window would have been B, the text output for the
        second window would be C, for the third window it would be A, for the fourth B, etc.
    window_duration : int
        The number of seconds to grab from each audio file before feature extraction.
    max_overlap_frac : float, optional
        Specify how much the windows overlap as a fraction of the window length.
        Default: 0.0.
    score_targets_csv : Optional[str]
        If present, filename of a CSV file with columns "uuid" and f"scores_{task}" for each task in
        `label_type`. These scores will be packaged with the quantized labels for the elements of
        this dataset, so they can be used by the model e.g. as targets for knowledge distillation.
        This CSV is typically generated by:
          (1) Running evaluation of a model using `eval_set = "train"`.
          (2) Identifying the WANDB_RUN_ID of this eval job
          (3) Running `kirad.utils.get_dataframe_with_scores(WANDB_RUN_ID, "test").to_csv(OUT_CSV)`
              (note this must be called with split="test", not split="train" -- see the comment
              in `get_dataframe_with_scores` for a description of this inconsistency.)
        but could alternatively be generated by some other ad hoc process, e.g. combining scores
        of multiple models in some way.
    """

    model_config = ConfigDict(
        extra="forbid",
        arbitrary_types_allowed=True,
    )

    # Specify input
    df_or_dataset_version: pd.DataFrame | DatasetVersion | str | Path | Mapping | list[
        pd.DataFrame | DatasetVersion | str | Path | Mapping
    ]
    data_sources: Optional[DataSource | list[DataSource]] = None
    dataset_splits: Optional[DatasetSplit | list[DatasetSplit]] = None
    dataset_fold: Optional[int] = None
    # Specify columns in the metadata
    label_column: Mapping[str, str] = DEFAULT_LABEL_COLUMN
    duration_column: Optional[str] = None
    # Preprocessing options
    merge_rows_by_subject: MERGE_ROWS_BY_SUBJECT_VALUES = False
    window_duration: int = INFERENCE_WINDOW_SIZE
    max_overlap_frac: float = 0.0
    # Filtering options
    fraction: float = 1.0
    min_duration: Optional[float] = INFERENCE_WINDOW_SIZE
    max_duration: Optional[float] = None
    # Label creation options
    score_targets_csv: Optional[str | Path] = None
    audio_root: Optional[str | Path] = None
    text_root: Optional[str | Path] = None

    # Private attributes that are not a part of the config
    _df: pd.DataFrame = PrivateAttr()

    def __init__(self, *args, **kwargs):
        """Load and concatenate all metadata sources, index by uuid, preprocess, filter."""
        super().__init__(*args, **kwargs)
        # Accept one metadata source or a list of them.
        if not isinstance(self.df_or_dataset_version, list):
            self.df_or_dataset_version = [self.df_or_dataset_version]
        # Load every source and concatenate into one table of recordings.
        self._df = pd.concat(
            load_relative_dataset(
                item,
                self.data_sources,
                self.dataset_splits,
                self.dataset_fold,
                audio_root=self.audio_root,
                text_root=self.text_root,
            )
            for item in self.df_or_dataset_version
        )
        # Index rows (samples) by UUID
        self._df.set_index("uuid", inplace=True)
        self.preprocess_metadata()
        self.filter_metadata()

    def preprocess_metadata(self):
        """Join KD score targets, optionally merge rows per subject, count windows per recording.

        Adds column num_windows (windows of the "all" method per recording).
        Note: uses its own max_overlap_frac; keep it equal to the audio dataset's.
        """
        # Join teacher (KD) scores by uuid as extra scores_<task> columns.
        if self.score_targets_csv is not None:
            score_target_df = pd.read_csv(self.score_targets_csv)
            score_target_df.set_index("uuid", inplace=True)
            self._df = self._df.merge(
                score_target_df,
                how="left",
                left_index=True,
                right_index=True,
                suffixes=(None, "_target"),
            )

        # Merge rows of the same subject (audio is concatenated across their rows).
        if self.merge_rows_by_subject:
            self._df = merge_df_rows_by_subject(self._df)

            if self.merge_rows_by_subject == "many_to_many":
                # Expand each subject to one row per text_index (0..n-1).
                self._df["num_files"] = self._df["filename"].map(
                    lambda fn: 1 if isinstance(fn, str) else len(fn)
                )
                max_files = self._df["num_files"].max()
                window_indices = pd.DataFrame(
                    [(i, j) for i in range(1, max_files + 1) for j in range(i)],
                    columns=["num_files", "text_index"],
                )
                self._df = self._df.merge(window_indices, on="num_files", how="left")

        # Calculate the maximum number of audio windows per sample based on window
        # duration and window overlap fraction.
        self._df["num_windows"] = self._df[self.duration_column].map(
            lambda d: len(
                get_window_starts(
                    "all",
                    int(d * EXPECTED_SAMPLE_RATE),
                    self.window_duration * EXPECTED_SAMPLE_RATE,
                    self.max_overlap_frac,
                )
            )
        )

    def filter_metadata(self):
        """Subsample `fraction` of rows, then drop recordings outside min/max duration (s)."""
        # HARDCODED: random_state=1 for the `fraction` subsample
        self._df = self._df.sample(frac=self.fraction, random_state=1)
        logger.info(f"Keeping {self.fraction:0.2%} of samples.")
        orig_len = len(self._df)
        mask = pd.Series([True for _ in range(0, orig_len)], index=self._df.index)

        # Filter out files based on audio duration.
        for duration, op_func, logging_str in [
            (self.min_duration, operator.ge, "less than"),
            (self.max_duration, operator.le, "greater than"),
        ]:
            if duration is not None:
                if self.duration_column is None:
                    raise ValueError(
                        "`duration_column` must be specified to filter metadata by "
                        "duration."
                    )
                duration_mask = op_func(self._df[self.duration_column], duration)
                logger.info(
                    f"Filtered out {(~duration_mask).sum()} files that were "
                    f"{logging_str} {duration} seconds."
                )
                mask &= duration_mask

        # Apply all filters
        self._df = self._df[mask]
        logger.info(
            f"""
            Original number of files\t: {orig_len}
            Remaining number of files\t: {mask.sum()}
            Number of files filtered\t: {(~mask).sum()}
            """
        )

    def __getitem__(self, uuid: str) -> dict[str, str | dict[str, int | float]]:
        """Return {uuid, label: {task: {primary: raw score, kd: teacher score}}} for one UUID.

        kd is added only when a scores_<task> column exists (from score_targets_csv).
        """
        row = self._df.loc[uuid]
        # Raw score of each task from its CSV label column (primary label).
        label = {
            task: {DatasetFields.LabelFields.PRIMARY: row[label_column]}
            for task, label_column in self.label_column.items()
        }

        # Add KD targets if provided
        for task in self.label_column:
            if f"scores_{task}" in row:
                label[task][DatasetFields.LabelFields.KD] = row[f"scores_{task}"]

        return {DatasetFields.UUID: uuid, DatasetFields.LABEL: label}

    def __len__(self) -> int:
        """Number of recordings after filtering."""
        return len(self._df)

model_config = ConfigDict(extra='forbid', arbitrary_types_allowed=True) class-attribute instance-attribute

df_or_dataset_version instance-attribute

data_sources = None class-attribute instance-attribute

dataset_splits = None class-attribute instance-attribute

dataset_fold = None class-attribute instance-attribute

label_column = DEFAULT_LABEL_COLUMN class-attribute instance-attribute

duration_column = None class-attribute instance-attribute

merge_rows_by_subject = False class-attribute instance-attribute

window_duration = INFERENCE_WINDOW_SIZE class-attribute instance-attribute

max_overlap_frac = 0.0 class-attribute instance-attribute

fraction = 1.0 class-attribute instance-attribute

min_duration = INFERENCE_WINDOW_SIZE class-attribute instance-attribute

max_duration = None class-attribute instance-attribute

score_targets_csv = None class-attribute instance-attribute

audio_root = None class-attribute instance-attribute

text_root = None class-attribute instance-attribute

__init__(*args, **kwargs)

Load and concatenate all metadata sources, index by uuid, preprocess, filter.

Source code in kirad/datasets_and_dataloaders.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
def __init__(self, *args, **kwargs):
    """Load and concatenate all metadata sources, index by uuid, preprocess, filter."""
    super().__init__(*args, **kwargs)
    # Accept one metadata source or a list of them.
    if not isinstance(self.df_or_dataset_version, list):
        self.df_or_dataset_version = [self.df_or_dataset_version]
    # Load every source and concatenate into one table of recordings.
    self._df = pd.concat(
        load_relative_dataset(
            item,
            self.data_sources,
            self.dataset_splits,
            self.dataset_fold,
            audio_root=self.audio_root,
            text_root=self.text_root,
        )
        for item in self.df_or_dataset_version
    )
    # Index rows (samples) by UUID
    self._df.set_index("uuid", inplace=True)
    self.preprocess_metadata()
    self.filter_metadata()

preprocess_metadata()

Join KD score targets, optionally merge rows per subject, count windows per recording.

Adds column num_windows (windows of the "all" method per recording). Note: uses its own max_overlap_frac; keep it equal to the audio dataset's.

Source code in kirad/datasets_and_dataloaders.py 
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
def preprocess_metadata(self):
    """Join KD score targets, optionally merge rows per subject, count windows per recording.

    Adds column num_windows (windows of the "all" method per recording).
    Note: uses its own max_overlap_frac; keep it equal to the audio dataset's.
    """
    # Join teacher (KD) scores by uuid as extra scores_<task> columns.
    if self.score_targets_csv is not None:
        score_target_df = pd.read_csv(self.score_targets_csv)
        score_target_df.set_index("uuid", inplace=True)
        self._df = self._df.merge(
            score_target_df,
            how="left",
            left_index=True,
            right_index=True,
            suffixes=(None, "_target"),
        )

    # Merge rows of the same subject (audio is concatenated across their rows).
    if self.merge_rows_by_subject:
        self._df = merge_df_rows_by_subject(self._df)

        if self.merge_rows_by_subject == "many_to_many":
            # Expand each subject to one row per text_index (0..n-1).
            self._df["num_files"] = self._df["filename"].map(
                lambda fn: 1 if isinstance(fn, str) else len(fn)
            )
            max_files = self._df["num_files"].max()
            window_indices = pd.DataFrame(
                [(i, j) for i in range(1, max_files + 1) for j in range(i)],
                columns=["num_files", "text_index"],
            )
            self._df = self._df.merge(window_indices, on="num_files", how="left")

    # Calculate the maximum number of audio windows per sample based on window
    # duration and window overlap fraction.
    self._df["num_windows"] = self._df[self.duration_column].map(
        lambda d: len(
            get_window_starts(
                "all",
                int(d * EXPECTED_SAMPLE_RATE),
                self.window_duration * EXPECTED_SAMPLE_RATE,
                self.max_overlap_frac,
            )
        )
    )

filter_metadata()

Subsample fraction of rows, then drop recordings outside min/max duration (s).

Source code in kirad/datasets_and_dataloaders.py 
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
def filter_metadata(self):
    """Subsample `fraction` of rows, then drop recordings outside min/max duration (s)."""
    # HARDCODED: random_state=1 for the `fraction` subsample
    self._df = self._df.sample(frac=self.fraction, random_state=1)
    logger.info(f"Keeping {self.fraction:0.2%} of samples.")
    orig_len = len(self._df)
    mask = pd.Series([True for _ in range(0, orig_len)], index=self._df.index)

    # Filter out files based on audio duration.
    for duration, op_func, logging_str in [
        (self.min_duration, operator.ge, "less than"),
        (self.max_duration, operator.le, "greater than"),
    ]:
        if duration is not None:
            if self.duration_column is None:
                raise ValueError(
                    "`duration_column` must be specified to filter metadata by "
                    "duration."
                )
            duration_mask = op_func(self._df[self.duration_column], duration)
            logger.info(
                f"Filtered out {(~duration_mask).sum()} files that were "
                f"{logging_str} {duration} seconds."
            )
            mask &= duration_mask

    # Apply all filters
    self._df = self._df[mask]
    logger.info(
        f"""
        Original number of files\t: {orig_len}
        Remaining number of files\t: {mask.sum()}
        Number of files filtered\t: {(~mask).sum()}
        """
    )

__getitem__(uuid)

Return {uuid, label: {task: {primary: raw score, kd: teacher score}}} for one UUID.

kd is added only when a scores_ column exists (from score_targets_csv).

Source code in kirad/datasets_and_dataloaders.py 
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
def __getitem__(self, uuid: str) -> dict[str, str | dict[str, int | float]]:
    """Return {uuid, label: {task: {primary: raw score, kd: teacher score}}} for one UUID.

    kd is added only when a scores_<task> column exists (from score_targets_csv).
    """
    row = self._df.loc[uuid]
    # Raw score of each task from its CSV label column (primary label).
    label = {
        task: {DatasetFields.LabelFields.PRIMARY: row[label_column]}
        for task, label_column in self.label_column.items()
    }

    # Add KD targets if provided
    for task in self.label_column:
        if f"scores_{task}" in row:
            label[task][DatasetFields.LabelFields.KD] = row[f"scores_{task}"]

    return {DatasetFields.UUID: uuid, DatasetFields.LABEL: label}

__len__()

Number of recordings after filtering.

Source code in kirad/datasets_and_dataloaders.py 
312
313
314
def __len__(self) -> int:
    """Number of recordings after filtering."""
    return len(self._df)

KintsugiDatasetBase

Bases: BaseModel, Dataset

Base for modality datasets: holds the shared KintsugiMetadata; length = its rows.

Source code in kirad/datasets_and_dataloaders.py 
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
class KintsugiDatasetBase(
    BaseModel,
    Dataset,
):
    """Base for modality datasets: holds the shared KintsugiMetadata; length = its rows."""

    model_config = ConfigDict(
        extra="forbid",
        arbitrary_types_allowed=True,
    )
    metadata: KintsugiMetadata

    def __len__(self) -> int:
        """Number of recordings in the shared metadata."""
        return len(self.metadata)

model_config = ConfigDict(extra='forbid', arbitrary_types_allowed=True) class-attribute instance-attribute

metadata instance-attribute

__len__()

Number of recordings in the shared metadata.

Source code in kirad/datasets_and_dataloaders.py 
329
330
331
def __len__(self) -> int:
    """Number of recordings in the shared metadata."""
    return len(self.metadata)

KintsugiAudioDataset

Bases: KintsugiDatasetBase

Audio by UUID: loads WAV, optional augmentations, cuts windows, optional log-mel.

Shapes: raw [W, 480000] if preprocessor=False, else HF log-mel [W, 80, 3000].

An audio dataset that processes and generates log-mel spectrogram features from audio files.

Attributes:

Name Type Description
metadata KintsugiMetadata

Metadata object
preprocessor string

Which preprocessor to use. Allowable values can be found at

https://huggingface.co/models?search=openai/whisper

and include "openai/whisper-large", "openai/whisper-medium.en", "openai/whisper-small.en" (default), or "openai/whisper-base.en".
normalize_audio (bool, optional)

If true, then scale the audio waveform to be between -1 and 1.
normalize_features (bool, optional)

If true, then normalize the mel-spectrogram to have zero-mean, unit-variance.
window_duration int

The number of seconds to grab from each audio file before feature extraction
window_method (str, optional)

Specify how that audio is windowed prior to extracting features. The window length is specified by window_duration. Can be "all" (default), "single_random", "single_start", or "random_adjacent_pair".
max_overlap_frac (float, optional)

Specify how much the windows overlap as a fraction of the window length. Default: 0.0.
ideal_logmel_energies (list or ndarray, optional)

Specifies the value of each log-mel bin to normalize to. The length of the array must match the number of log-mel bins (eg. 80 for WhisperFeatureExtractor).
Source code in kirad/datasets_and_dataloaders.py 
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
class KintsugiAudioDataset(KintsugiDatasetBase):
    """Audio by UUID: loads WAV, optional augmentations, cuts windows, optional log-mel.

    Shapes: raw [W, 480000] if preprocessor=False, else HF log-mel [W, 80, 3000].

    An audio dataset that processes and generates log-mel spectrogram features from
    audio files.

    Attributes
    ----------
    metadata : KintsugiMetadata
        Metadata object
    preprocessor : string
        Which preprocessor to use. Allowable values can be found at

            https://huggingface.co/models?search=openai/whisper

        and include "openai/whisper-large", "openai/whisper-medium.en",
        "openai/whisper-small.en" (default), or "openai/whisper-base.en".
    normalize_audio : bool, optional
        If true, then scale the audio waveform to be between -1 and 1.
    normalize_features : bool, optional
        If true, then normalize the mel-spectrogram to have zero-mean,
        unit-variance.
    window_duration : int
        The number of seconds to grab from each audio file before feature
        extraction
    window_method : str, optional
        Specify how that audio is windowed prior to extracting features. The window
        length is specified by `window_duration`. Can be "all" (default),
        "single_random", "single_start", or "random_adjacent_pair".
    max_overlap_frac : float, optional
        Specify how much the windows overlap as a fraction of the window length.
        Default: 0.0.
    ideal_logmel_energies : list or numpy.ndarray, optional
        Specifies the value of each log-mel bin to normalize to. The length of the
        array must match the number of log-mel bins (eg. 80 for
        WhisperFeatureExtractor).
    """

    # HARDCODED: HF feature extractor name; mar_2024/config.yaml sets false (log-mel on GPU)
    preprocessor: str | Literal[False] = "openai/whisper-small.en"
    normalize_audio: bool = False
    normalize_features: bool = False
    window_duration: int = INFERENCE_WINDOW_SIZE
    window_method: WINDOW_METHOD_VALUES = "all"
    max_overlap_frac: float = 0.0
    ideal_logmel_energies: Optional[npt.ArrayLike | str | Path] = None
    pad_last_chunk_to_full: bool = False
    augmentations: Sequence[Mapping] = ()

    # Private attributes that are not a part of the config
    _preprocessor_with_audio_normalization: Callable[..., torch.Tensor] = PrivateAttr()
    _augmentations: list[Tuple[float, Callable]] = PrivateAttr(default_factory=list)

    def __init__(self, *args, **kwargs):
        """Load ideal_logmel_energies, build augmentations from config, build preprocessor.

        Each augmentation config is {class_name, params, aug_prob}.
        """
        super().__init__(*args, **kwargs)
        # Load ideal log-mel bin energies (file or array) as a tensor for rescaling.
        if self.ideal_logmel_energies is not None:
            if isinstance(self.ideal_logmel_energies, str | Path):
                self.ideal_logmel_energies = np.load(self.ideal_logmel_energies)
            self.ideal_logmel_energies = torch.tensor(self.ideal_logmel_energies)

        # Instantiate each configured augmentation, paired with its apply probability.
        for augmentation_config in self.augmentations:
            class_name = augmentation_config["class_name"]
            params = augmentation_config["params"]
            aug_prob = augmentation_config["aug_prob"]

            aug_factory = import_object(class_name)
            augmentation = aug_factory(**params)

            self._augmentations.append((aug_prob, augmentation))

        self.init_data_processor()

    def init_data_processor(self):
        """Create the windowing + feature function (HF extractor or raw passthrough)."""
        # HF feature extractor, or raw passthrough when preprocessor is False.
        if self.preprocessor:
            preprocessor = AutoFeatureExtractor.from_pretrained(self.preprocessor)
        else:
            preprocessor = dummy_preprocessor
        self._preprocessor_with_audio_normalization = (
            get_preprocessor_with_audio_normalization(
                preprocessor=preprocessor,
                window_method=self.window_method,
                normalize_audio=self.normalize_audio,
                normalize_features=self.normalize_features,
                trim_duration=self.window_duration,
                ideal_logmel_energies=self.ideal_logmel_energies,
                max_overlap_frac=self.max_overlap_frac,
                pad_last_chunk_to_full=self.pad_last_chunk_to_full,
            )
        )

    def __getitem__(self, uuid: str) -> torch.Tensor:
        """Load, concatenate (channel 0), augment and window one recording's audio.

        Note: a sample rate other than 16 kHz only logs a warning; audio is not resampled.
        Shapes: [W, 480000] raw or [W, 80, 3000] log-mel.

        Loads an audio tensor.

        Parameters
        ----------
        uuid : str
            The UUID of the sample.

        Returns
        -------
        Tensor containing audio features.

        Raises
        ------
        RuntimeError
            If the sample rate of the audio file to be loaded is not
            `TARGET_SAMPLE_RATE`
        """
        # Fetch the row and list its audio files (merged subjects have several).
        row = self.metadata._df.loc[uuid]
        filenames = (
            row.audio_filename
            if isinstance(row.audio_filename, list)
            else [row.audio_filename]
        )
        # Load each file, keep channel 0 and concatenate into one waveform.
        audio = torch.zeros(0)
        srs = set()
        for fp in filenames:
            try:
                new_audio, sr = torchaudio.load(fp)
            except soundfile.LibsndfileError as e:
                raise ValueError(f"Failed to load audio from {fp}") from e

            if sr != EXPECTED_SAMPLE_RATE:
                logger.warning(
                    f"Found sample rate of {sr} for file {fp}, which doesn't match "
                    f"the expected sample rate of {EXPECTED_SAMPLE_RATE}."
                )

            if new_audio.dim() == 2:
                # HARDCODED: only channel 0 is used (stereo files lose the other channel)
                new_audio = new_audio[0]
            audio = torch.cat((audio, new_audio), dim=0)
            srs.add(sr)

        # All concatenated files must share one sample rate.
        if len(srs) != 1:
            raise ValueError(
                "Cannot concatenate audios with different sample rates: found sample "
                f"rates {srs} among file names {filenames}."
            )

        # applying augmentations
        for aug_prob, aug in self._augmentations:
            if random.random() < aug_prob:
                audio = aug(audio)

        # Cut windows and extract features: [W, ...] per recording.
        features = self._preprocessor_with_audio_normalization(audio, sr)
        # "window_index" is not created in this repo; branch keeps two adjacent windows
        if "window_index" in self.metadata._df.columns:
            window_index_a = row.window_index
            window_index_b = (window_index_a + 1) % row.num_windows
            features = torch.stack(
                [features[window_index_a, ...], features[window_index_b, ...]], dim=0
            )

        return features

preprocessor = 'openai/whisper-small.en' class-attribute instance-attribute

normalize_audio = False class-attribute instance-attribute

normalize_features = False class-attribute instance-attribute

window_duration = INFERENCE_WINDOW_SIZE class-attribute instance-attribute

window_method = 'all' class-attribute instance-attribute

max_overlap_frac = 0.0 class-attribute instance-attribute

ideal_logmel_energies = None class-attribute instance-attribute

pad_last_chunk_to_full = False class-attribute instance-attribute

augmentations = () class-attribute instance-attribute

__init__(*args, **kwargs)

Load ideal_logmel_energies, build augmentations from config, build preprocessor.

Each augmentation config is {class_name, params, aug_prob}.

Source code in kirad/datasets_and_dataloaders.py 
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def __init__(self, *args, **kwargs):
    """Load ideal_logmel_energies, build augmentations from config, build preprocessor.

    Each augmentation config is {class_name, params, aug_prob}.
    """
    super().__init__(*args, **kwargs)
    # Load ideal log-mel bin energies (file or array) as a tensor for rescaling.
    if self.ideal_logmel_energies is not None:
        if isinstance(self.ideal_logmel_energies, str | Path):
            self.ideal_logmel_energies = np.load(self.ideal_logmel_energies)
        self.ideal_logmel_energies = torch.tensor(self.ideal_logmel_energies)

    # Instantiate each configured augmentation, paired with its apply probability.
    for augmentation_config in self.augmentations:
        class_name = augmentation_config["class_name"]
        params = augmentation_config["params"]
        aug_prob = augmentation_config["aug_prob"]

        aug_factory = import_object(class_name)
        augmentation = aug_factory(**params)

        self._augmentations.append((aug_prob, augmentation))

    self.init_data_processor()

init_data_processor()

Create the windowing + feature function (HF extractor or raw passthrough).

Source code in kirad/datasets_and_dataloaders.py 
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
def init_data_processor(self):
    """Create the windowing + feature function (HF extractor or raw passthrough)."""
    # HF feature extractor, or raw passthrough when preprocessor is False.
    if self.preprocessor:
        preprocessor = AutoFeatureExtractor.from_pretrained(self.preprocessor)
    else:
        preprocessor = dummy_preprocessor
    self._preprocessor_with_audio_normalization = (
        get_preprocessor_with_audio_normalization(
            preprocessor=preprocessor,
            window_method=self.window_method,
            normalize_audio=self.normalize_audio,
            normalize_features=self.normalize_features,
            trim_duration=self.window_duration,
            ideal_logmel_energies=self.ideal_logmel_energies,
            max_overlap_frac=self.max_overlap_frac,
            pad_last_chunk_to_full=self.pad_last_chunk_to_full,
        )
    )

__getitem__(uuid)

Load, concatenate (channel 0), augment and window one recording's audio.

Note: a sample rate other than 16 kHz only logs a warning; audio is not resampled. Shapes: [W, 480000] raw or [W, 80, 3000] log-mel.

Loads an audio tensor.

Parameters:

Name Type Description Default
uuid str

The UUID of the sample.

 required

Returns:

Type Description
Tensor containing audio features.

Raises:

Type Description
RuntimeError

If the sample rate of the audio file to be loaded is not TARGET_SAMPLE_RATE
Source code in kirad/datasets_and_dataloaders.py 
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
def __getitem__(self, uuid: str) -> torch.Tensor:
    """Load, concatenate (channel 0), augment and window one recording's audio.

    Note: a sample rate other than 16 kHz only logs a warning; audio is not resampled.
    Shapes: [W, 480000] raw or [W, 80, 3000] log-mel.

    Loads an audio tensor.

    Parameters
    ----------
    uuid : str
        The UUID of the sample.

    Returns
    -------
    Tensor containing audio features.

    Raises
    ------
    RuntimeError
        If the sample rate of the audio file to be loaded is not
        `TARGET_SAMPLE_RATE`
    """
    # Fetch the row and list its audio files (merged subjects have several).
    row = self.metadata._df.loc[uuid]
    filenames = (
        row.audio_filename
        if isinstance(row.audio_filename, list)
        else [row.audio_filename]
    )
    # Load each file, keep channel 0 and concatenate into one waveform.
    audio = torch.zeros(0)
    srs = set()
    for fp in filenames:
        try:
            new_audio, sr = torchaudio.load(fp)
        except soundfile.LibsndfileError as e:
            raise ValueError(f"Failed to load audio from {fp}") from e

        if sr != EXPECTED_SAMPLE_RATE:
            logger.warning(
                f"Found sample rate of {sr} for file {fp}, which doesn't match "
                f"the expected sample rate of {EXPECTED_SAMPLE_RATE}."
            )

        if new_audio.dim() == 2:
            # HARDCODED: only channel 0 is used (stereo files lose the other channel)
            new_audio = new_audio[0]
        audio = torch.cat((audio, new_audio), dim=0)
        srs.add(sr)

    # All concatenated files must share one sample rate.
    if len(srs) != 1:
        raise ValueError(
            "Cannot concatenate audios with different sample rates: found sample "
            f"rates {srs} among file names {filenames}."
        )

    # applying augmentations
    for aug_prob, aug in self._augmentations:
        if random.random() < aug_prob:
            audio = aug(audio)

    # Cut windows and extract features: [W, ...] per recording.
    features = self._preprocessor_with_audio_normalization(audio, sr)
    # "window_index" is not created in this repo; branch keeps two adjacent windows
    if "window_index" in self.metadata._df.columns:
        window_index_a = row.window_index
        window_index_b = (window_index_a + 1) % row.num_windows
        features = torch.stack(
            [features[window_index_a, ...], features[window_index_b, ...]], dim=0
        )

    return features

KintsugiTextDataset

Bases: KintsugiDatasetBase

Transcript by UUID: reads ASR JSON and tokenizes it (BERT by default).

A text dataset that processes and tokenizes strings using a tokenizer (eg. BERT tokenizer).

Attributes:

Name Type Description
metadata KintsugiMetadata

Metadata object
preprocessor str

Which Tokenizer preprocessor to use. Can be "google-bert/bert-base-uncased". See

https://huggingface.co/models?search=google-bert/bert

for other possible values.
prompt str

Prompt for generative LLM model
max_tokens int

Max sequence length to which the sequence is padded or truncated at. If not provided, the default is used.
text_output_per_sample str

Specifies how to handle multiple text strings per sample (eg. due to setting merge_rows_by_subject in KintsugiMetadata). - "cyclic": starting from a specified index in the "text_index" column or a random index if "text_index" is missing, take num_audio_window strings from the string list (wrap around as needed, hence the name "cyclic") - "merge": merge all strings in the list into a single string.
Source code in kirad/datasets_and_dataloaders.py 
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
class KintsugiTextDataset(KintsugiDatasetBase):
    """Transcript by UUID: reads ASR JSON and tokenizes it (BERT by default).

    A text dataset that processes and tokenizes strings using a tokenizer (eg. BERT
    tokenizer).

    Attributes
    ----------
    metadata : KintsugiMetadata
        Metadata object
    preprocessor : str
        Which Tokenizer preprocessor to use. Can be "google-bert/bert-base-uncased".  See

            https://huggingface.co/models?search=google-bert/bert

        for other possible values.
    prompt : str
        Prompt for generative LLM model
    max_tokens : int
        Max sequence length to which the sequence is padded or truncated at.
        If not provided, the default is used.
    text_output_per_sample : str
        Specifies how to handle multiple text strings per sample (eg. due to setting
        `merge_rows_by_subject` in `KintsugiMetadata`).
        - "cyclic": starting from a specified index in the "text_index" column or a
          random index if "text_index" is missing, take num_audio_window strings from
          the string list (wrap around as needed, hence the name "cyclic")
        - "merge": merge all strings in the list into a single string.
    """

    # HARDCODED: tokenizer name (DAM 2 text branch); a W&B artifact URL is also accepted
    preprocessor: str | Path = "google-bert/bert-base-uncased"
    text_output_per_sample: TEXT_OUTPUT_PER_SAMPLE_VALUES = "merge"
    prompt: Optional[str] = None
    max_tokens: Optional[int] = None

    # Private attributes that are not a part of the config.
    _tokenizer: Callable[[str], BatchEncoding] = PrivateAttr()

    def __init__(self, *args, **kwargs):
        """Build the tokenizer after pydantic validation."""
        super().__init__(*args, **kwargs)
        self.init_data_processor()

    def init_data_processor(self):
        """Load the tokenizer (local/HF name or W&B artifact) and wrap it with get_tokenizer."""

        # Tokenizer location: download if W&B artifact URL, else name/path as is.
        tokenizer_path = self.preprocessor
        if self.preprocessor.startswith("https:"):
            run, name, tag = get_wandb_run_name_tag_from_artifact_url(self.preprocessor)
            tokenizer_path = localize_wandb_artifact(run, name, tag, return_dir=True)

        # use_fast=False prefers regular tokenizer with full functionality
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=False)
        self._tokenizer = get_tokenizer(tokenizer, self.prompt, self.max_tokens)

    def __getitem__(self, uuid: str) -> dict[str, torch.Tensor]:
        """Read the transcript JSON(s) of one recording and tokenize them.

        Note: only whisper_timestamped JSON ("text" field) works; the Google ASR
        fallback reopens the same path, so it never succeeds.
        Shapes: input_ids/attention_mask [1, L] ("merge") or [W, L] ("cyclic").

        Loads an audio tensor.

        Parameters
        ----------
        uuid : str
            The UUID of the sample.

        Returns
        -------
        Dictionary containing text features.
        """
        row = self.metadata._df.loc[uuid]

        # Experiments at https://app.clickup.com/20576956/v/dc/kkynw-19251/kkynw-112271
        # showed that whisper large transcripts perform better than google transcripts
        filenames = (
            row.text_filename
            if isinstance(row.text_filename, list)
            else [row.text_filename]
        )
        asr_output_ids = (
            row.asr_output_id
            if isinstance(row.asr_output_id, list)
            else [row.asr_output_id]
        )
        # Read each transcript (Whisper JSON; Google fallback reopens the same path).
        texts = []
        for json_filename, asr_output_id in zip(filenames, asr_output_ids):
            try:
                texts.append(read_whisper_timestamped_json(json_filename))
            except FileNotFoundError:
                texts.append(read_google_asr_json(json_filename))

        # cyclic: one text per audio window; merge: join all texts into a single string.
        if self.text_output_per_sample == "cyclic":
            num_audio_windows = row["num_windows"]
            text_index = row.get("text_index", random.randint(0, len(texts) - 1))
            text_output = [
                texts[(text_index + i) % len(texts)] for i in range(num_audio_windows)
            ]
        elif self.text_output_per_sample == "merge":
            text_output = " ".join(texts)

        # Tokenize the chosen text(s) into input_ids / attention_mask.
        features = self._tokenizer(text_output)

        return features

preprocessor = 'google-bert/bert-base-uncased' class-attribute instance-attribute

text_output_per_sample = 'merge' class-attribute instance-attribute

prompt = None class-attribute instance-attribute

max_tokens = None class-attribute instance-attribute

__init__(*args, **kwargs)

Build the tokenizer after pydantic validation.

Source code in kirad/datasets_and_dataloaders.py 
554
555
556
557
def __init__(self, *args, **kwargs):
    """Build the tokenizer after pydantic validation."""
    super().__init__(*args, **kwargs)
    self.init_data_processor()

init_data_processor()

Load the tokenizer (local/HF name or W&B artifact) and wrap it with get_tokenizer.

Source code in kirad/datasets_and_dataloaders.py 
559
560
561
562
563
564
565
566
567
568
569
570
def init_data_processor(self):
    """Load the tokenizer (local/HF name or W&B artifact) and wrap it with get_tokenizer."""

    # Tokenizer location: download if W&B artifact URL, else name/path as is.
    tokenizer_path = self.preprocessor
    if self.preprocessor.startswith("https:"):
        run, name, tag = get_wandb_run_name_tag_from_artifact_url(self.preprocessor)
        tokenizer_path = localize_wandb_artifact(run, name, tag, return_dir=True)

    # use_fast=False prefers regular tokenizer with full functionality
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=False)
    self._tokenizer = get_tokenizer(tokenizer, self.prompt, self.max_tokens)

__getitem__(uuid)

Read the transcript JSON(s) of one recording and tokenize them.

Note: only whisper_timestamped JSON ("text" field) works; the Google ASR fallback reopens the same path, so it never succeeds. Shapes: input_ids/attention_mask [1, L] ("merge") or [W, L] ("cyclic").

Loads an audio tensor.

Parameters:

Name Type Description Default
uuid str

The UUID of the sample.

 required

Returns:

Type Description
Dictionary containing text features.
Source code in kirad/datasets_and_dataloaders.py 
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
def __getitem__(self, uuid: str) -> dict[str, torch.Tensor]:
    """Read the transcript JSON(s) of one recording and tokenize them.

    Note: only whisper_timestamped JSON ("text" field) works; the Google ASR
    fallback reopens the same path, so it never succeeds.
    Shapes: input_ids/attention_mask [1, L] ("merge") or [W, L] ("cyclic").

    Loads an audio tensor.

    Parameters
    ----------
    uuid : str
        The UUID of the sample.

    Returns
    -------
    Dictionary containing text features.
    """
    row = self.metadata._df.loc[uuid]

    # Experiments at https://app.clickup.com/20576956/v/dc/kkynw-19251/kkynw-112271
    # showed that whisper large transcripts perform better than google transcripts
    filenames = (
        row.text_filename
        if isinstance(row.text_filename, list)
        else [row.text_filename]
    )
    asr_output_ids = (
        row.asr_output_id
        if isinstance(row.asr_output_id, list)
        else [row.asr_output_id]
    )
    # Read each transcript (Whisper JSON; Google fallback reopens the same path).
    texts = []
    for json_filename, asr_output_id in zip(filenames, asr_output_ids):
        try:
            texts.append(read_whisper_timestamped_json(json_filename))
        except FileNotFoundError:
            texts.append(read_google_asr_json(json_filename))

    # cyclic: one text per audio window; merge: join all texts into a single string.
    if self.text_output_per_sample == "cyclic":
        num_audio_windows = row["num_windows"]
        text_index = row.get("text_index", random.randint(0, len(texts) - 1))
        text_output = [
            texts[(text_index + i) % len(texts)] for i in range(num_audio_windows)
        ]
    elif self.text_output_per_sample == "merge":
        text_output = " ".join(texts)

    # Tokenize the chosen text(s) into input_ids / attention_mask.
    features = self._tokenizer(text_output)

    return features

KintsugiCacheDataset

Bases: KintsugiDatasetBase

Precomputed backbone features by UUID (DAM 3 teacher cache), optional dim slice.

A dataset that returns features from a backbone cache.

Attributes:

Name Type Description
metadata KintsugiMetadata

Metadata object
backbone_cache pandas DataFrame with 2 columns [uuid, feat]
Sometimes(e.g. LLMA training case) we can be interested in a specific feature slice of the
feature vectors in the backbone cache.
feat_start_idx and feat_end_idx params define start and end indices of such a slice.
feat_end_idx is a non-inclusive index.
Source code in kirad/datasets_and_dataloaders.py 
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
class KintsugiCacheDataset(KintsugiDatasetBase):
    """Precomputed backbone features by UUID (DAM 3 teacher cache), optional dim slice.

    A dataset that returns features from a backbone cache.

    Attributes
    ----------
    metadata : KintsugiMetadata
        Metadata object
    backbone_cache : pandas DataFrame with 2 columns [uuid, feat]

    Sometimes(e.g. LLMA training case) we can be interested in a specific feature slice of the
    feature vectors in the backbone cache.
    feat_start_idx and feat_end_idx params define start and end indices of such a slice.
    feat_end_idx is a non-inclusive index.
    """

    backbone_cache: pd.DataFrame | str | Path
    feat_start_idx: Optional[int] = 0
    feat_end_idx: Optional[int] = None

    def __init__(self, *args, **kwargs):
        """Load the cache (DataFrame or W&B artifact URL) and index it by uuid."""
        super().__init__(*args, **kwargs)
        if isinstance(self.backbone_cache, str | Path):
            # Load cache from W&B
            cache_run, cache_name, cache_tag = get_wandb_run_name_tag_from_artifact_url(
                self.backbone_cache
            )
            cache_path = localize_wandb_artifact(cache_run, cache_name, cache_tag)
            state_dict = torch.load(cache_path)
            self.backbone_cache = pd.DataFrame(
                state_dict["backbone_cache"].items(), columns=["uuid", "feat"]
            )
        # Cache must expose uuid and feat columns.
        if not {"uuid", "feat"}.issubset(set(self.backbone_cache.columns)):
            raise ValueError(
                "Backbone cache data frame should have columns 'uuid' and 'feat'."
            )

        # Index the cache by UUIDs
        if (
            self.backbone_cache.index.name is None
            or self.backbone_cache.index.name != "uuid"
        ):
            if self.backbone_cache.index.name != "uuid":
                self.backbone_cache.reset_index(inplace=True)
            self.backbone_cache.set_index("uuid", inplace=True)

        self._process_cache()

    def _process_cache(self):
        """Drop metadata rows whose UUID is not in the cache (modifies shared metadata)."""

        # Filter out UUIDs from the metadata that aren't in the backbone cache.
        mask = self.metadata._df.index.isin(self.backbone_cache.index)
        logger.info(
            f"Filtered out {(~mask).sum()} files that are not in the backbone cache."
        )
        if self.metadata.fraction < 1:
            logger.info(
                f"Warning: 'fraction' parameter applies to the original metadata, "
                f"and it may have an approximate effect with respect to the cache content."
            )
        self.metadata._df = self.metadata._df[mask]

    def __getitem__(self, uuid: str) -> torch.Tensor:
        """Cached features of one UUID, sliced to [feat_start_idx:feat_end_idx] on dim 1."""
        # we assume the features are contained in dim = 1, but tensor can have
        # an arbitrary number of dimensions
        return self.backbone_cache.loc[uuid]["feat"][
            :, self.feat_start_idx : self.feat_end_idx, ...
        ]

backbone_cache instance-attribute

feat_start_idx = 0 class-attribute instance-attribute

feat_end_idx = None class-attribute instance-attribute

__init__(*args, **kwargs)

Load the cache (DataFrame or W&B artifact URL) and index it by uuid.

Source code in kirad/datasets_and_dataloaders.py 
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
def __init__(self, *args, **kwargs):
    """Load the cache (DataFrame or W&B artifact URL) and index it by uuid."""
    super().__init__(*args, **kwargs)
    if isinstance(self.backbone_cache, str | Path):
        # Load cache from W&B
        cache_run, cache_name, cache_tag = get_wandb_run_name_tag_from_artifact_url(
            self.backbone_cache
        )
        cache_path = localize_wandb_artifact(cache_run, cache_name, cache_tag)
        state_dict = torch.load(cache_path)
        self.backbone_cache = pd.DataFrame(
            state_dict["backbone_cache"].items(), columns=["uuid", "feat"]
        )
    # Cache must expose uuid and feat columns.
    if not {"uuid", "feat"}.issubset(set(self.backbone_cache.columns)):
        raise ValueError(
            "Backbone cache data frame should have columns 'uuid' and 'feat'."
        )

    # Index the cache by UUIDs
    if (
        self.backbone_cache.index.name is None
        or self.backbone_cache.index.name != "uuid"
    ):
        if self.backbone_cache.index.name != "uuid":
            self.backbone_cache.reset_index(inplace=True)
        self.backbone_cache.set_index("uuid", inplace=True)

    self._process_cache()

__getitem__(uuid)

Cached features of one UUID, sliced to [feat_start_idx:feat_end_idx] on dim 1.

Source code in kirad/datasets_and_dataloaders.py 
694
695
696
697
698
699
700
def __getitem__(self, uuid: str) -> torch.Tensor:
    """Cached features of one UUID, sliced to [feat_start_idx:feat_end_idx] on dim 1."""
    # we assume the features are contained in dim = 1, but tensor can have
    # an arbitrary number of dimensions
    return self.backbone_cache.loc[uuid]["feat"][
        :, self.feat_start_idx : self.feat_end_idx, ...
    ]

UuidSampler

Bases: Sampler[str]

Yields metadata UUIDs (shuffled for train, sequential for val/test) for the loader.

Provides a random sampler to sample UUIDs from the metadata.

Source code in kirad/datasets_and_dataloaders.py 
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
class UuidSampler(Sampler[str]):
    """Yields metadata UUIDs (shuffled for train, sequential for val/test) for the loader.

    Provides a random sampler to sample UUIDs from the metadata.
    """

    def __init__(
        self,
        metadata: KintsugiMetadata,
        shuffle: bool = False,
        distributed: bool = False,
    ):
        """Wrap a Random, Sequential or Distributed sampler over the UUID list."""
        self.uuids = metadata._df.index.tolist()
        # Pick order: distributed shard, random permutation (train) or sequential.
        if distributed:
            self.sampler = DistributedSampler(self.uuids, shuffle=shuffle)
        else:
            if shuffle:
                self.sampler = RandomSampler(
                    self.uuids, replacement=False, num_samples=None, generator=None
                )
            else:
                self.sampler = SequentialSampler(self.uuids)

    def __iter__(self) -> Iterator[str]:
        """Yield UUIDs in sampler order."""
        for idx in self.sampler:
            yield self.uuids[idx]

    def __len__(self) -> int:
        """Number of UUIDs."""
        return len(self.uuids)

uuids = metadata._df.index.tolist() instance-attribute

sampler = DistributedSampler(self.uuids, shuffle=shuffle) instance-attribute

__init__(metadata, shuffle=False, distributed=False)

Wrap a Random, Sequential or Distributed sampler over the UUID list.

Source code in kirad/datasets_and_dataloaders.py 
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
def __init__(
    self,
    metadata: KintsugiMetadata,
    shuffle: bool = False,
    distributed: bool = False,
):
    """Wrap a Random, Sequential or Distributed sampler over the UUID list."""
    self.uuids = metadata._df.index.tolist()
    # Pick order: distributed shard, random permutation (train) or sequential.
    if distributed:
        self.sampler = DistributedSampler(self.uuids, shuffle=shuffle)
    else:
        if shuffle:
            self.sampler = RandomSampler(
                self.uuids, replacement=False, num_samples=None, generator=None
            )
        else:
            self.sampler = SequentialSampler(self.uuids)

__iter__()

Yield UUIDs in sampler order.

Source code in kirad/datasets_and_dataloaders.py 
728
729
730
731
def __iter__(self) -> Iterator[str]:
    """Yield UUIDs in sampler order."""
    for idx in self.sampler:
        yield self.uuids[idx]

__len__()

Number of UUIDs.

Source code in kirad/datasets_and_dataloaders.py 
733
734
735
def __len__(self) -> int:
    """Number of UUIDs."""
    return len(self.uuids)

dummy_preprocessor(x, *args, **kwargs)

Stack raw window arrays into a tensor [W, samples]; used when preprocessor=False.

Source code in kirad/datasets_and_dataloaders.py 
334
335
336
def dummy_preprocessor(x, *args, **kwargs):
    """Stack raw window arrays into a tensor [W, samples]; used when preprocessor=False."""
    return torch.from_numpy(np.array(x))

get_window_starts(method, audio_len, inference_window_samples, max_inference_window_overlap_fraction=MAX_INFERENCE_WINDOW_OVERLAP_FRACTION)

Start sample(s) of the windows cut from one recording, by window_method.

"all": evenly spaced via np.linspace, hop >= (1 - overlap) * window, >= 1 window. "single_start": 0. "single_random": one random start. "random_adjacent_pair": 2 starts.

Calculate the audio windows based on the audio duration. Each audio window is inference_window_samples long.

A maximum overlap of max_inference_window_overlap_samples is used, but the overlap is flexible so that as much audio as possible is used.

Arguments 
method : The method used for selecting window starts
audio_len : the length of the audio in samples
inference_window_samples : the length of the inference window in samples
max_inference_window_overlap_fraction : fraction of each inference window allowed to overlap with the previous
    one (between 0.0 and 1.0)

Returns:

Type Description
For method == "all", a list of window starts overlapping by at most the specified fraction.

For method == "single_random", a single window start at a random position.
Source code in kirad/datasets_and_dataloaders.py 
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
def get_window_starts(
    method: WINDOW_METHOD_VALUES,
    audio_len: int,
    inference_window_samples: int,
    max_inference_window_overlap_fraction: float = MAX_INFERENCE_WINDOW_OVERLAP_FRACTION,
) -> int | Iterable[int]:
    """Start sample(s) of the windows cut from one recording, by window_method.

    "all": evenly spaced via np.linspace, hop >= (1 - overlap) * window, >= 1 window.
    "single_start": 0. "single_random": one random start. "random_adjacent_pair": 2 starts.

    Calculate the audio windows based on the audio duration. Each audio window is inference_window_samples long.

    A maximum overlap of max_inference_window_overlap_samples is used, but the overlap is flexible so that as much audio
    as possible is used.

    Arguments
    ---------
        method : The method used for selecting window starts
        audio_len : the length of the audio in samples
        inference_window_samples : the length of the inference window in samples
        max_inference_window_overlap_fraction : fraction of each inference window allowed to overlap with the previous
            one (between 0.0 and 1.0)

    Returns
    -------
        For method == "all", a list of window starts overlapping by at most the specified fraction.
        For method == "single_random", a single window start at a random position.

    """
    overflow_len = audio_len - inference_window_samples
    # Fixed-position and random-position methods return start(s) directly.
    if method == "single_start":
        return 0
    if method == "single_random":
        return random.randint(0, max(overflow_len, 0))
    if method == "random_adjacent_pair":
        if audio_len < 2 * inference_window_samples:  # overlap unavoidable
            return [0, overflow_len]  # minimize overlap
        first = random.randint(0, audio_len - 2 * inference_window_samples)
        return [first, first + inference_window_samples]
    # "all": evenly spaced starts, hop >= (1 - max overlap) * window; at least 1 window.
    if method == "all":
        min_hop_samples = int(
            (1 - max_inference_window_overlap_fraction) * inference_window_samples
        )

        n_windows = 1 + overflow_len // min_hop_samples
        return np.linspace(0, overflow_len, max(n_windows, 1)).astype(int)
    raise ValueError(f"Unknown window method: {method}")

get_preprocessor_with_audio_normalization(preprocessor, window_method, normalize_audio=False, normalize_features=False, trim_duration=INFERENCE_WINDOW_SIZE, ideal_logmel_energies=None, max_overlap_frac=MAX_INFERENCE_WINDOW_OVERLAP_FRACTION, pad_last_chunk_to_full=False)

Build the function audio [T] -> windows -> features used by KintsugiAudioDataset.

Steps: optional DC removal + peak scaling to [-1, 1], optional zero-pad to full windows, cut windows (get_window_starts), run preprocessor, optional per-bin shift so each log-mel bin's time-mean equals ideal_logmel_energies. Shapes: [W, window samples] raw, or [W, 80, 3000] with the HF Whisper extractor.

Source code in kirad/datasets_and_dataloaders.py 
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
def get_preprocessor_with_audio_normalization(
    preprocessor: Callable,
    window_method: WINDOW_METHOD_VALUES,
    normalize_audio: bool = False,
    normalize_features: bool = False,
    trim_duration: int = INFERENCE_WINDOW_SIZE,
    ideal_logmel_energies: Optional[torch.FloatTensor] = None,
    max_overlap_frac: float = MAX_INFERENCE_WINDOW_OVERLAP_FRACTION,
    pad_last_chunk_to_full: bool = False,
) -> Callable[..., torch.Tensor]:
    """Build the function audio [T] -> windows -> features used by KintsugiAudioDataset.

    Steps: optional DC removal + peak scaling to [-1, 1], optional zero-pad to full
    windows, cut windows (get_window_starts), run preprocessor, optional per-bin
    shift so each log-mel bin's time-mean equals ideal_logmel_energies.
    Shapes: [W, window samples] raw, or [W, 80, 3000] with the HF Whisper extractor.
    """

    def forward(
        audio: torch.Tensor, sampling_rate: int = EXPECTED_SAMPLE_RATE
    ) -> torch.Tensor:
        """Turn one recording's waveform [T] into stacked window features [W, ...]."""
        if normalize_audio:
            # Remove DC offset and scale amplitude to [-1, 1]
            audio = audio - torch.mean(audio)
            audio = audio / torch.max(torch.abs(audio))

        # window length uses the file's own sample rate (30 s * 16 kHz = 480,000)
        chunk_samples = sampling_rate * trim_duration

        # pad audio in a way, so that the last chunk is not dropped
        if pad_last_chunk_to_full:
            if max_overlap_frac > 0:
                raise ValueError(
                    f"pad_last_chunk_to_full is only supported for non-overlapping windows"
                )
            num_chunks = np.ceil(len(audio) / chunk_samples)
            pad_size = int(num_chunks * chunk_samples - len(audio))
            audio = torch.nn.functional.pad(audio, (0, pad_size))

        # Window start sample(s) for this recording.
        window_starts = get_window_starts(
            window_method, len(audio), chunk_samples, max_overlap_frac
        )

        # Wrap a bare int start (single_start / single_random) into a list.
        if isinstance(window_starts, Iterable):
            all_window_starts = window_starts
        else:
            all_window_starts = [window_starts]
        # Slice each window from the waveform and run the preprocessor (e.g. log-mel).
        features = preprocessor(
            [
                audio[start : start + chunk_samples].numpy(force=True)
                for start in all_window_starts
            ],
            return_tensors="pt",
            sampling_rate=sampling_rate,
            do_normalize=normalize_features,
        )
        # Unwrap the HF output object to the feature tensor.
        for key in ("input_features", "input_values"):
            if hasattr(features, key):
                features = getattr(features, key)
                break

        # Shift each log-mel bin so its time-mean equals the ideal energy.
        if ideal_logmel_energies is not None:
            mean_features = torch.mean(features, dim=-1)
            # features are [batch, n_logmel_bins, n_frames]
            rescale_factor = ideal_logmel_energies.unsqueeze(0) - mean_features
            rescale_factor = rescale_factor.unsqueeze(2)
            features += rescale_factor
        return features

    return forward

get_tokenizer(tokenizer, prompt=None, max_length=None)

Wrap a HF tokenizer: optional chat prompt, pad/truncate to max_length, return pt tensors.

Get a callable function that tokenizes text with a specified tokenizer.

Parameters:

Name Type Description Default
tokenizer PreTrainedTokenizerBase

The tokenizer to invoke (eg. BertTokenizer.from_pretrained("bert-base-uncased"))

 required
prompt str

Text which prompts generative LLM model. If provided, assume we are dealing with generative LLM.

 None
max_length Optional[int]

Max sequence length. If provided, it will override default max sequence length.

 None

Returns:

Type Description
Callable

Function that takes text and returns the tokenization in a dictionary containing the following: - "input_ids": token IDs corresponding to the tokenization of the text - "token_type_ids": token type IDs - "attention_mask": specifies which tokens the LLM should attend to
Source code in kirad/datasets_and_dataloaders.py 
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
def get_tokenizer(
    tokenizer: PreTrainedTokenizerBase,
    prompt: Optional[str] = None,
    max_length: Optional[int] = None,
) -> Callable[[str], BatchEncoding]:
    """Wrap a HF tokenizer: optional chat prompt, pad/truncate to max_length, return pt tensors.

    Get a callable function that tokenizes text with a specified tokenizer.

    Parameters
    ----------
    tokenizer : transformers.PreTrainedTokenizerBase
        The tokenizer to invoke (eg. BertTokenizer.from_pretrained("bert-base-uncased"))

    prompt : str
        Text which prompts generative LLM model. If provided, assume we are dealing with generative LLM.

    max_length: int
        Max sequence length. If provided, it will override default max sequence length.

    Returns
    -------
    Callable
        Function that takes text and returns the tokenization in a dictionary
        containing the following:
            - "input_ids": token IDs corresponding to the tokenization of the text
            - "token_type_ids": token type IDs
            - "attention_mask": specifies which tokens the LLM should attend to
    """

    # Not every tokenizer has pad token assigned by default.
    # For some (such as Llama) it should be specified manually.
    if tokenizer.pad_token is None:
        if tokenizer.eos_token:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "[PAD]"})

    def forward(text: str) -> BatchEncoding:
        """Tokenize one string (or list of strings) padded to max_length."""
        # If a prompt is specified assume we are dealing with generative LLM model
        if prompt is not None:
            messages = [
                {"role": "system", "content": prompt},
                {"role": "user", "content": text},
            ]

            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )

        # Some tokenizers don't have "prepare_for_tokenization" method
        if hasattr(tokenizer, "prepare_for_tokenization"):
            text, _ = tokenizer.prepare_for_tokenization(text)

        # Tokenize, padded/truncated to max_length (model max if None).
        tokenized = tokenizer(
            text,
            add_special_tokens=True,
            max_length=tokenizer.model_max_length if max_length is None else max_length,
            padding="max_length",
            truncation=True,
            return_token_type_ids=True,
            return_attention_mask=True,
            return_tensors="pt",
        )
        return tokenized

    return forward