Files, function by function

One card per file, in story order. Step links to the walkthrough; Calls and Called by come from a static analysis of the code (approximate: dynamic calls such as Lightning hooks and import_object are not seen).

scripts/run_whisper_asr.py

Offline transcription of one WAV into the transcript JSON used by text models. · open file ↗

Function	Step	Calls	Called by	What it does
run_whisper_asr		—	—	Transcribe a 16 kHz WAV (optional start/end in s) with Whisper large-v3 to JSON.

scripts/launcher.py

Command-line entry point: python scripts/launcher.py <mode> config.yaml. · open file ↗

Function	Step	Calls	Called by	What it does
launcher	1	_log_datasets_to_wandb, _log_to_wandb, _setup_callbacks, _setup_launcher, run_caching, run_compilation, run_evaluation, run_training	—	Run one mode (train/evaluate/cache/compile) end to end from a YAML config.
run_training	12	LauncherConfig.get_data_loader launcher_utils.py	launcher	Fit on train/val, reload the best checkpoint, then test it on the test split.
run_evaluation		LauncherConfig.get_data_loader launcher_utils.py	launcher	Test a trained checkpoint (model_params.ckpt_path) on the test split.
_setup_launcher	2	LauncherConfig launcher_utils.py, LauncherLogger launcher_utils.py	launcher	Load the YAML into LauncherConfig and create the W&B logger on rank 0 only.
_setup_callbacks	10	DistributedModelCheckpoint launcher_utils.py	launcher	Best-checkpoint saver, early stopping and LR monitor on "mode:metric" monitor_loss.
_log_to_wandb	11	LauncherLogger.log_code launcher_utils.py, LauncherLogger.log_table_artifacts launcher_utils.py, get_package_info utils.py, get_package_path utils.py	launcher	Upload experiment + kirad source, package versions and ideal log-mel energies to W&B.
_log_datasets_to_wandb	43	LauncherLogger.log_table_artifacts launcher_utils.py	launcher	Upload the metadata table actually used for each split as a W&B "dataset" artifact.
run_caching		LauncherConfig.get_data_loader_kwargs launcher_utils.py, LauncherConfig.get_dataset launcher_utils.py, UuidSampler datasets_and_dataloaders.py, cache_backbone utils.py	launcher	Run the backbone over train/val/test and store its outputs in W&B (cache mode).
run_compilation		LauncherConfig.get_data_loader_kwargs launcher_utils.py, LauncherConfig.get_dataset launcher_utils.py	launcher	Export the model with torch.export on one test window, check outputs, log .pt2 to W&B.
_parse_args		—	—	Parse CLI args: mode, config path, optional --experiment-dir.
setup_multi_node_environment		—	—	Wait for all Kubernetes job pods, then set MASTER_ADDR/PORT, WORLD_SIZE, NODE_RANK.

kirad/launcher_utils.py

Turns the YAML config into objects: params, model, datasets, loaders, W&B logger. · open file ↗

Function	Step	Calls	Called by	What it does
DistributedModelCheckpoint		—	_setup_callbacks launcher.py	ModelCheckpoint that also keeps a copy of the best checkpoint on every node.
DistributedModelCheckpoint.on_validation_end	42	—	—	Save best/last ckpt after validation (with tuned thresholds); replicate on nodes.
SchedulerParams		—	—	LR scheduler settings from training_params.scheduler (gamma, milestones).
OptimizerParams		—	—	Optimizer settings for one parameter group (lr, weight_decay).
TrainingParams	4	—	LauncherConfig.__init__	Validated training_params: batch sizes, epochs, monitor metric, optimizer, seed.
TrainingParams.__init__		—	—	Compute the global batch size (WORLD_SIZE x batch_size) and seed everything.
TrainingParams.total_accum_grad_batches		—	—	Gradient-accumulation steps: effective_batch_size // global batch size.
ModelParams		—	LauncherConfig.__init__	Validated model_params: model class path, config kwargs and checkpoint to load.
ModelParams.__init__		get_wandb_run_name_tag_from_artifact_url utils.py, localize_wandb_artifact utils.py	—	Resolve ckpt_path (W&B download or backbone cache) and check restored config.
LauncherConfig		—	_setup_launcher launcher.py	Everything built from one YAML config; datasets and loaders are created lazily.
LauncherConfig.__init__	3	LauncherConfig._init_compute_params, LauncherConfig._init_data_params, LauncherConfig._init_model, ModelParams, TrainingParams	—	Read YAML, parse training/data/model params and instantiate the model.
LauncherConfig._init_compute_params		—	LauncherConfig.__init__	Count GPUs/CPUs, set DataLoader workers per GPU and deterministic torch/CuBLAS.
LauncherConfig._init_data_params	5	LauncherConfig.get_params_by_split, import_object utils.py	LauncherConfig.__init__	Parse data_params keys "<modality>=<class.path>" into factories + per-split params.
LauncherConfig._init_model	7	import_object utils.py	LauncherConfig.__init__	Import model_path, build the model and load ckpt weights (or defer full restore).
LauncherConfig.get_metadata		—	LauncherConfig.get_data_loader_kwargs, LauncherConfig.get_dataset	Build (once) and return the metadata object for a split.
LauncherConfig.get_dataset	14	LauncherConfig.get_metadata	LauncherConfig.get_data_loader, run_caching launcher.py, run_compilation launcher.py	StackDataset of metadata + each modality dataset for a split (built once).
LauncherConfig.get_data_loader_kwargs		LauncherConfig.get_metadata, UuidSampler datasets_and_dataloaders.py	LauncherConfig.get_data_loader, run_caching launcher.py, run_compilation launcher.py	DataLoader kwargs: broadcast_collate, UuidSampler (shuffled + drop_last if train).
LauncherConfig.get_data_loader	13	LauncherConfig.get_data_loader_kwargs, LauncherConfig.get_dataset	run_evaluation launcher.py, run_training launcher.py	DataLoader for a split (cached per split and mode); batch_size = recordings.
LauncherConfig.resolve_artifact_path		—	LauncherConfig.resolve_artifact_paths_in_param_dict	Resolve the absolute path for a file path specified in the config file.
LauncherConfig.resolve_artifact_paths_in_param_dict		LauncherConfig.resolve_artifact_path	LauncherConfig.get_params_by_split	Make relative string params absolute when they point to an existing file.
LauncherConfig.get_params_by_split	6	LauncherConfig.resolve_artifact_paths_in_param_dict	LauncherConfig._init_data_params	Merge config params for one split: default < val_test (val/test) < split.
LauncherLogger		—	_setup_launcher launcher.py	W&B run for the experiment; logs code and pandas tables as artifacts.
LauncherLogger.__init__		—	—	Start a WandbLogger (log_model=True uploads checkpoints) under log_path/wandb.
LauncherLogger.log_code		—	_log_to_wandb launcher.py	Upload a source folder as the code artifact <artifact_name>-<run_id>.
LauncherLogger.log_table_artifacts		dataframe_to_wandb_table utils.py	_log_datasets_to_wandb launcher.py, _log_to_wandb launcher.py	Upload {artifact: {table_name: DataFrame}} as W&B table artifacts.

kirad/utils.py

Shared helpers: per-recording averaging, W&B artifact I/O, dynamic imports, analysis. · open file ↗

Function	Step	Calls	Called by	What it does
load_audio		—	—	Load channel 0 of an audio file, optionally resampling; not used by the datasets.
average_tensor_in_segments		—	DepAnxClassifier.forward_backbone mtl_aug_2024/model.py, MultiTaskOrdinalRegressionPLModule.predict_step base_models.py, OrdinalRegressionPLModule.compute_and_log_loss base_models.py, OrdinalRegressionPLModule.predict_step base_models.py, WhisperLLMA.forward mtl_aug_2024/model.py, WhisperLLMA.forward_ground_truth mtl_aug_2024/model.py, average_and_variance_tensor_in_segments	Mean over each recording's windows: [ΣW, ...] -> [B, ...] using length.
average_and_variance_tensor_in_segments	34	average_tensor_in_segments	OrdinalRegressionPLModule.compute_and_log_loss base_models.py, OrdinalRegressionPLModule.step base_models.py	Per-recording mean and (biased) variance of window scores: [ΣW, ...] -> two [B, ...].
localize_wandb_artifact		—	DepAnxClassifier.__init__ mtl_aug_2024/model.py, KintsugiCacheDataset.__init__ datasets_and_dataloaders.py, KintsugiTextDataset.init_data_processor datasets_and_dataloaders.py, LLAMABackbone.__init__ base_models.py, ModelParams.__init__ launcher_utils.py, combine_backbone_modality_state_dicts_from_wandb, get_dataframe_with_scores, get_wandb_report_card_tsvs	Download a W&B artifact once to ~/.cache/kintsugi and return its local path.
get_wandb_run_name_tag_from_artifact_url		—	DepAnxClassifier.__init__ mtl_aug_2024/model.py, KintsugiCacheDataset.__init__ datasets_and_dataloaders.py, KintsugiTextDataset.init_data_processor datasets_and_dataloaders.py, LLAMABackbone.__init__ base_models.py, ModelParams.__init__ launcher_utils.py	Parse a W&B artifact URL into (run that logged it, artifact name, tag).
TSVAccumulator		—	get_wandb_report_card_tsvs	Helper for formatting tab-separated-value data with section headers in the first row.
TSVAccumulator.__init__		—	—	Start with an empty header and one empty row.
TSVAccumulator.set_header		—	get_wandb_report_card_tsvs	Set the section prefix used by later write calls.
TSVAccumulator.write		—	get_wandb_report_card_tsvs	Add one column: "<header> - <first>" in row 0, then rest down the rows.
TSVAccumulator.__str__		—	—	Render rows as TSV text.
get_wandb_report_card_tsvs		TSVAccumulator, TSVAccumulator.set_header, TSVAccumulator.write, indet_analysis, localize_wandb_artifact	get_wandb_report_card_tsvs_with_instructions	Build a TSV report card per task from a W&B run and log it back as an artifact.
get_wandb_report_card_tsvs_with_instructions		get_wandb_report_card_tsvs	—	Report cards from get_wandb_report_card_tsvs, prefixed with paste instructions.
str_to_tensor		—	—	Encode a python string as 1-d torch.Tensor of uint8s.
strs_to_tensor		—	OrdinalRegressionPLModule.step base_models.py	Encode strings as a zero-padded 2-d uint8 tensor (uuids for MultiCatMetric).
tensor_to_str		—	—	Decode a 1-d torch.Tensor of uint8s into a python string (strips zero padding).
tensor_to_strs		—	OrdinalRegressionPLModule.on_validation_test_epoch_end base_models.py	Decode a 2-d torch.Tensor of uint8s into a list of python strings.
get_dataframe_with_scores	44	load_relative_dataset dataset_utils.py, localize_wandb_artifact	indet_analysis	Join a run's metadata CSVs with its W&B scores-<task>-<split> tables on uuid.
indet_analysis		IndetRocCurve.auc metrics.py, IndetRocCurve.sn_eq_sp metrics.py, IndetSnSpArray.roc_curve metrics.py, get_dataframe_with_scores	get_wandb_report_card_tsvs	Tune sn_eq_sp thresholds on val per budget and binary cut, then measure them on test.
analyze_dataframe_with_scores		IndetRocCurve.auc metrics.py, IndetRocCurve.sn_eq_sp metrics.py	—	Mean AUROC and sn_eq_sp (0% indet) over binary cuts, per subgroup of the dataframe.
get_package_path		—	_log_to_wandb launcher.py, get_package_info	Repo root of kipy or kirad (env KIPY_PATH / KINTSUGI_RAD_PATH, else from import).
get_package_info		get_package_path	_log_to_wandb launcher.py	Version (from setup.cfg) and git commit/dirty state of kipy or kirad.
import_object		—	KintsugiAudioDataset.__init__ datasets_and_dataloaders.py, LauncherConfig._init_data_params launcher_utils.py, LauncherConfig._init_model launcher_utils.py	Import a class/function from a dotted path, e.g. "model.DepAnxClassifier".
transfer_to_cuda		—	cache_backbone	Move a tensor or a (nested) dict of tensors to a CUDA device.
cache_backbone		transfer_to_cuda	run_caching launcher.py	Run the backbone on all datasets, gather outputs by uuid across ranks, log to W&B.
read_whisper_timestamped_json		—	KintsugiTextDataset.__getitem__ datasets_and_dataloaders.py	Transcript text ("text" field) from a whisper-timestamped JSON.
read_google_asr_json		—	KintsugiTextDataset.__getitem__ datasets_and_dataloaders.py	Transcript from a Google Speech API JSON: top alternative of each result, joined.
flat_to_nested_dict		—	combine_backbone_modality_state_dicts	Split dot-separated keys into nested dicts.
nested_to_flat_dict		—	combine_backbone_modality_state_dicts	Flatten nested dicts into dot-separated keys.
combine_backbone_modality_state_dicts		flat_to_nested_dict, nested_to_flat_dict	combine_backbone_modality_state_dicts_from_wandb	Merge backbones of several checkpoints into one DAM 2 state_dict (backbone.<key>.*).
combine_backbone_modality_state_dicts_from_wandb		combine_backbone_modality_state_dicts, localize_wandb_artifact	—	Save a merged DAM 2 backbone checkpoint built from the "best" models of W&B runs.
dataframe_to_wandb_table		—	LauncherLogger.log_table_artifacts launcher_utils.py, OrdinalRegressionPLModule.log_test_indet_analysis base_models.py, OrdinalRegressionPLModule.on_validation_test_epoch_end base_models.py	wandb.Table from a DataFrame, raising W&B's row limit so large tables are not cut.
open_single_file_wandb_artifact		—	wandb_table_artifact_to_dataframe	Download the only file of a W&B artifact and return an open handle.
wandb_table_artifact_to_dataframe		open_single_file_wandb_artifact	—	DataFrame from a single-table W&B artifact, read directly from its JSON.

kirad/datasets_and_dataloaders.py

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

Function	Step	Calls	Called by	What it does
KintsugiMetadata	15	—	—	Metadata table indexed by UUID: loads CSV, merges KD scores, filters by duration.
KintsugiMetadata.__init__		KintsugiMetadata.filter_metadata, KintsugiMetadata.preprocess_metadata, load_relative_dataset dataset_utils.py	—	Load and concatenate all metadata sources, index by uuid, preprocess, filter.
KintsugiMetadata.preprocess_metadata	17	get_window_starts, merge_df_rows_by_subject dataset_utils.py	KintsugiMetadata.__init__	Join KD score targets, optionally merge rows per subject, count windows per recording.
KintsugiMetadata.filter_metadata	18	—	KintsugiMetadata.__init__	Subsample fraction of rows, then drop recordings outside min/max duration (s).
KintsugiMetadata.__getitem__	21	—	—	Return {uuid, label: {task: {primary: raw score, kd: teacher score}}} for one UUID.
KintsugiMetadata.__len__		—	—	Number of recordings after filtering.
KintsugiDatasetBase		—	—	Base for modality datasets: holds the shared KintsugiMetadata; length = its rows.
KintsugiDatasetBase.__len__		—	—	Number of recordings in the shared metadata.
dummy_preprocessor		—	—	Stack raw window arrays into a tensor [W, samples]; used when preprocessor=False.
KintsugiAudioDataset		—	—	Audio by UUID: loads WAV, optional augmentations, cuts windows, optional log-mel.
KintsugiAudioDataset.__init__		KintsugiAudioDataset.init_data_processor, import_object utils.py	—	Load ideal_logmel_energies, build augmentations from config, build preprocessor.
KintsugiAudioDataset.init_data_processor		get_preprocessor_with_audio_normalization	KintsugiAudioDataset.__init__	Create the windowing + feature function (HF extractor or raw passthrough).
KintsugiAudioDataset.__getitem__	22	—	—	Load, concatenate (channel 0), augment and window one recording's audio.
KintsugiTextDataset		—	—	Transcript by UUID: reads ASR JSON and tokenizes it (BERT by default).
KintsugiTextDataset.__init__		KintsugiTextDataset.init_data_processor	—	Build the tokenizer after pydantic validation.
KintsugiTextDataset.init_data_processor		get_tokenizer, get_wandb_run_name_tag_from_artifact_url utils.py, localize_wandb_artifact utils.py	KintsugiTextDataset.__init__	Load the tokenizer (local/HF name or W&B artifact) and wrap it with get_tokenizer.
KintsugiTextDataset.__getitem__		read_google_asr_json utils.py, read_whisper_timestamped_json utils.py	—	Read the transcript JSON(s) of one recording and tokenize them.
KintsugiCacheDataset		—	—	Precomputed backbone features by UUID (DAM 3 teacher cache), optional dim slice.
KintsugiCacheDataset.__init__		KintsugiCacheDataset._process_cache, get_wandb_run_name_tag_from_artifact_url utils.py, localize_wandb_artifact utils.py	—	Load the cache (DataFrame or W&B artifact URL) and index it by uuid.
KintsugiCacheDataset._process_cache		—	KintsugiCacheDataset.__init__	Drop metadata rows whose UUID is not in the cache (modifies shared metadata).
KintsugiCacheDataset.__getitem__		—	—	Cached features of one UUID, sliced to [feat_start_idx:feat_end_idx] on dim 1.
UuidSampler	20	—	LauncherConfig.get_data_loader_kwargs launcher_utils.py, run_caching launcher.py	Yields metadata UUIDs (shuffled for train, sequential for val/test) for the loader.
UuidSampler.__init__		—	—	Wrap a Random, Sequential or Distributed sampler over the UUID list.
UuidSampler.__iter__		—	—	Yield UUIDs in sampler order.
UuidSampler.__len__		—	—	Number of UUIDs.
get_window_starts	23	—	KintsugiMetadata.preprocess_metadata, get_preprocessor_with_audio_normalization	Start sample(s) of the windows cut from one recording, by window_method.
get_preprocessor_with_audio_normalization	19	get_window_starts	KintsugiAudioDataset.init_data_processor	Build the function audio [T] -> windows -> features used by KintsugiAudioDataset.
get_tokenizer		—	KintsugiTextDataset.init_data_processor	Wrap a HF tokenizer: optional chat prompt, pad/truncate to max_length, return pt tensors.

kirad/dataset_utils.py

Data stage helpers: load metadata CSVs, build file paths, merge rows, collate batches. · open file ↗

Function	Step	Calls	Called by	What it does
load_dataset		load_dataset_version	load_relative_dataset	Return metadata as a DataFrame: pass-through, read a CSV path, or a Kintsugi version.
load_relative_dataset	16	load_dataset	KintsugiMetadata.__init__ datasets_and_dataloaders.py, get_dataframe_with_scores utils.py	Load metadata and add audio_filename / text_filename paths from the filename column.
load_dataset_version		—	load_dataset	Fetch a Kintsugi dataset version (gsutil rsync) and read its split/fold CSVs.
resolve_audio_path		resolve_audio_path_from_bucket	—	Return a local audio directory: a given path, or a synced feature-store bucket.
resolve_audio_path_from_bucket		get_cache_dir	resolve_audio_path	Download a feature-store audio bucket with gsutil (if missing) and return its path.
get_cache_dir		—	resolve_audio_path_from_bucket	Local directory for a bucket name (AUDIO_DIR/<bucket>) or a given path.
calculate_distributions		get_data_type	split_dataset	Label and column distributions (histograms / category masses) used by split_dataset.
split_dataset		calculate_distributions	—	Speaker-disjoint train/test folds, optionally matching given distributions (by KL).
get_data_type		—	calculate_distributions	Return "numerical" if the Series converts to numbers, else "categorical".
merge_df_rows_by_subject		—	KintsugiMetadata.preprocess_metadata datasets_and_dataloaders.py	Merge rows sharing user + labels + demographics into one row per subject.
repeat_along_batch_dim		—	broadcast_along_batch_dim	Tile a tensor (or dict of tensors) repetitions times along dim 0.
broadcast_along_batch_dim		repeat_along_batch_dim	broadcast_collate	Tile every tensor of one sample to the largest dim 0 (= its number of windows W).
get_collate_fn		—	—	torch default collate, but tensors are merged with tensor_collate_fn (e.g. concat).
broadcast_collate	24	broadcast_along_batch_dim	—	DataLoader collate: concat windows of all recordings; length = windows per recording.

kirad/augmentations.py

Data stage (optional): RawBoost waveform augmentation, applied in KintsugiAudioDataset. · open file ↗

Function	Step	Calls	Called by	What it does
randRange		—	ISD_additive_noise, SSI_additive_noise, genNotchCoeffs	Uniform random value in [x1, x2), cast to int if integer.
normWav		—	ISD_additive_noise, LnL_convolutive_noise, SSI_additive_noise	Peak-normalize to [-1, 1] always, or only when the peak exceeds 1.
genNotchCoeffs		randRange	LnL_convolutive_noise, SSI_additive_noise	Random FIR filter: cascade of nBands random notch (band-stop) filters, random gain.
filterFIR		—	LnL_convolutive_noise, SSI_additive_noise	Apply FIR filter b to x, compensating the filter delay (same output length).
LnL_convolutive_noise		filterFIR, genNotchCoeffs, normWav	RawBoost.__call__	LnL: sum of random FIR-filtered powers x, x^2, ..., x^N_f (channel distortion).
ISD_additive_noise		normWav, randRange	RawBoost.__call__	ISD: perturb up to P% of random samples by signal-proportional impulsive noise.
SSI_additive_noise		filterFIR, genNotchCoeffs, normWav, randRange	RawBoost.__call__	SSI: add coloured Gaussian noise at a random SNR in [SNRmin, SNRmax] dB.
RawBoost		—	—	RawBoost waveform augmentation (LnL/ISD/SSI noise); default algo 4 = all three.
RawBoost.__init__		—	—	Store RawBoost parameters (paper defaults); raises if algo > 8.
RawBoost.__call__		ISD_additive_noise, LnL_convolutive_noise, SSI_additive_noise	—	Augment one waveform [T]; returns a tensor with the input dtype and device.

kirad/constants.py

Global constants: sample rate, score ranges, label cutoffs, batch field names. · open file ↗

Function	Step	Calls	Called by	What it does
DatasetFields		—	—	Key names of a batch dict: uuid, features, label.<task>.primary/kd, length, loss.

research/stable/ordinal_regression_mtl_mar_2024/model.py

DAM 1 model (stage: model): audio-only depression/anxiety classifier. · open file ↗

Function	Step	Calls	Called by	What it does
DepAnxClassifier	8	—	—	DAM 1: log-mel -> Whisper small.en encoder (LoRA) -> mean pool -> multitask CORAL head.
DepAnxClassifier.__init__		MultitaskHead base_models.py, OrdinalRegressionPLModule base_models.py, WhisperBackbone base_models.py, WhisperTorchFeatureExtractor feature_extraction_whisper.py	—	Build per-task CORAL modules, GPU log-mel extractor, Whisper backbone and head.
DepAnxClassifier.forward	27	DepAnxClassifier.forward_backbone	—	Backbone then MultitaskHead -> ({task: [ΣW, 1] scores}, lengths).
DepAnxClassifier.forward_backbone	28	—	DepAnxClassifier.compute_features_to_cache, DepAnxClassifier.forward	Pad raw audio to 480,000 samples, log-mel on GPU, Whisper encoder -> [ΣW, 768].
DepAnxClassifier.compute_features_to_cache		DepAnxClassifier.forward_backbone	—	Backbone embeddings split per recording (tuple of [W_i, 768]) for cache.
DepAnxClassifier.configure_optimizers		—	—	AdamW with two groups (backbone, head + CORAL biases) and MultiStepLR.

kirad/feature_extraction_whisper.py

Model stage: raw audio windows -> Whisper log-mel features, computed on the GPU in torch. · open file ↗

Function	Step	Calls	Called by	What it does
WhisperTorchFeatureExtractor		—	DepAnxClassifier.__init__ mtl_mar_2024/model.py	Torch Whisper log-mel extractor: 80 bins, 25 ms window, 10 ms hop, 16 kHz.
WhisperTorchFeatureExtractor.__init__		—	—	Precompute Hann window and Slaney mel filters; register the energy profile.
WhisperTorchFeatureExtractor._torch_extract_fbank_features		—	WhisperTorchFeatureExtractor.forward	Log10 mel spectrogram, clipped to 8 below the max, scaled (x + 4) / 4 as Whisper.
WhisperTorchFeatureExtractor.forward	29	WhisperTorchFeatureExtractor._torch_extract_fbank_features	—	Normalize each window, compute log-mel, shift bins to IDEAL_LOGMEL_ENERGIES.

kirad/base_models.py

Model building blocks. · open file ↗

Function	Step	Calls	Called by	What it does
OrdinalRegressionPLModule		—	DepAnxClassifier.__init__ mtl_aug_2024/model.py, DepAnxClassifier.__init__ mtl_mar_2024/model.py, create_task_modules_from_params mtl_jun_2024/model.py	An ordinal regression model with CORAL loss and threshold tuning.
OrdinalRegressionPLModule.__init__		CORALLoss losses.py, CORALThresholding ordinal_thresholding.py, IndetSnSpTensorModule metrics.py, MaxAccuracyOrdinalThresholding ordinal_thresholding.py, MaxMacroF1OrdinalThresholding ordinal_thresholding.py, MaxMacroPrecisionOrdinalThresholding ordinal_thresholding.py, MaxMacroRecallOrdinalThresholding ordinal_thresholding.py, MinAbsoluteErrorOrdinalThresholding ordinal_thresholding.py, MultiCatMetric metrics.py	—	Build CORAL loss, six threshold tuners, Sn=Sp buffers and val/test accumulators.
OrdinalRegressionPLModule.load_state_dict		—	—	Resize indet-threshold buffers to the checkpoint's shapes, then load weights.
OrdinalRegressionPLModule.forward		—	OrdinalRegressionPLModule.predict_step, OrdinalRegressionPLModule.step	Identity: the multi-task parent computes scores; this module only adds loss/metrics.
OrdinalRegressionPLModule.quantize_labels	38	—	OrdinalRegressionPLModule.on_validation_test_epoch_end	Map input labels (e.g. PHQ-9 sum) to output labels (e.g. 0/1/2 for L/M/H).
OrdinalRegressionPLModule.log_ordinal_metrics	41	IndetRocCurve.auc metrics.py, IndetRocCurve.sn_eq_sp metrics.py, IndetSnEqSpGraph metrics.py, IndetSnEqSpGraph.at_budget metrics.py, IndetSnSpArray.roc_curve metrics.py	OrdinalRegressionPLModule.on_validation_test_epoch_end	Log AUROC and min(Sn, Sp) per binary cut and indet budget, plus their mean.
OrdinalRegressionPLModule.log_thresholding_metrics		OptimalOrdinalThresholdingViaDynamicProgramming.best_constant_output_classifier ordinal_thresholding.py, OrdinalThresholding.is_valid ordinal_thresholding.py	OrdinalRegressionPLModule.on_validation_test_epoch_end	Log each tuned method's cost when scored with every method's thresholds.
OrdinalRegressionPLModule.log_test_indet_analysis		dataframe_to_wandb_table utils.py	OrdinalRegressionPLModule.on_validation_test_epoch_end	W&B line plots of test Sn, Sp, indet fraction vs the val indeterminate budget.
OrdinalRegressionPLModule.on_validation_epoch_end		OrdinalRegressionPLModule.on_validation_test_epoch_end	—	Lightning hook: run the val epoch-end logic and free CUDA cache.
OrdinalRegressionPLModule.on_test_epoch_end		OrdinalRegressionPLModule.on_validation_test_epoch_end	—	Lightning hook: run the test epoch-end logic.
OrdinalRegressionPLModule.on_validation_test_epoch_end	37	IndetSnSpArray.sn_eq_sp_graph metrics.py, OrdinalRegressionPLModule.log_ordinal_metrics, OrdinalRegressionPLModule.log_test_indet_analysis, OrdinalRegressionPLModule.log_thresholding_metrics, OrdinalRegressionPLModule.quantize_labels, dataframe_to_wandb_table utils.py, tensor_to_strs utils.py	OrdinalRegressionPLModule.on_test_epoch_end, OrdinalRegressionPLModule.on_validation_epoch_end	Gather scores, log W&B table, tune thresholds on val, compute Sn/Sp metrics.
OrdinalRegressionPLModule.compute_and_log_loss	33	average_and_variance_tensor_in_segments utils.py, average_tensor_in_segments utils.py	OrdinalRegressionPLModule.step	Loss = coral_w·CORAL + var_w·Var(window scores) + kd_w·MSE(teacher), logged.
OrdinalRegressionPLModule.step	32	OrdinalRegressionPLModule.compute_and_log_loss, OrdinalRegressionPLModule.forward, average_and_variance_tensor_in_segments utils.py, strs_to_tensor utils.py	OrdinalRegressionPLModule.test_step, OrdinalRegressionPLModule.training_step, OrdinalRegressionPLModule.validation_step	Accumulate per-recording mean/var scores for val/test, then return the loss.
OrdinalRegressionPLModule.training_step		OrdinalRegressionPLModule.step	—	Lightning train step for a single-task subclass.
OrdinalRegressionPLModule.validation_step		OrdinalRegressionPLModule.step	—	Lightning val step; loss is logged, not returned.
OrdinalRegressionPLModule.test_step		OrdinalRegressionPLModule.step	—	Lightning test step; loss is logged, not returned.
OrdinalRegressionPLModule.predict_step		OrdinalRegressionPLModule.forward, average_tensor_in_segments utils.py	—	Per-recording score: forward, then mean over its windows. Shapes: [B, 1].
MultiTaskOrdinalRegressionPLModule		—	—	Multi-task wrapper: one OrdinalRegressionPLModule per task on a shared model.
MultiTaskOrdinalRegressionPLModule.__init__		—	—	Store the per-task modules (loss, thresholds, accumulators) in a ModuleDict.
MultiTaskOrdinalRegressionPLModule.load_state_dict		—	—	Resize every task's indet-threshold buffers, then load weights.
MultiTaskOrdinalRegressionPLModule.configure_model		MultiTaskOrdinalRegressionPLModule.get_sub_logger	—	Give each task the trainer and a logger that prefixes names with '<task>/'.
MultiTaskOrdinalRegressionPLModule.get_sub_logger		—	MultiTaskOrdinalRegressionPLModule.configure_model	Return a log function that prefixes metric names with '<key>/'.
MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end	36	—	MultiTaskOrdinalRegressionPLModule.on_test_epoch_end, MultiTaskOrdinalRegressionPLModule.on_validation_epoch_end	Run each task's epoch end, then log harmonic mean of shared metrics across tasks.
MultiTaskOrdinalRegressionPLModule.on_validation_epoch_end		MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end	—	Lightning hook: val epoch end for all tasks.
MultiTaskOrdinalRegressionPLModule.on_test_epoch_end		MultiTaskOrdinalRegressionPLModule.on_validation_test_epoch_end	—	Lightning hook: test epoch end for all tasks.
MultiTaskOrdinalRegressionPLModule.compute_features_to_cache		—	—	Backbone embeddings split per recording, for the cache command (abstract).
MultiTaskOrdinalRegressionPLModule.training_step	25	MultiTaskOrdinalRegressionPLModule.step	—	Lightning train step: returns the step dict with total_loss.
MultiTaskOrdinalRegressionPLModule.validation_step		MultiTaskOrdinalRegressionPLModule.step	—	Lightning val step; losses are logged, not returned.
MultiTaskOrdinalRegressionPLModule.test_step		MultiTaskOrdinalRegressionPLModule.step	—	Lightning test step; losses are logged, not returned.
MultiTaskOrdinalRegressionPLModule.step	26	—	MultiTaskOrdinalRegressionPLModule.test_step, MultiTaskOrdinalRegressionPLModule.training_step, MultiTaskOrdinalRegressionPLModule.validation_step	Run model once, sum each task's loss (task.step) into total_loss.
MultiTaskOrdinalRegressionPLModule.predict_step		average_tensor_in_segments utils.py	—	Per-recording scores for each task: {task: [B, 1]} (mean over windows).
MultiTaskOrdinalRegressionBackboneDictPLModule		—	—	Multi-task module whose backbone is a ModuleDict; loads DAM 1 weights into "audio".
MultiTaskOrdinalRegressionBackboneDictPLModule.load_state_dict		—	—	Rename backbone.base_model.* keys to backbone.audio.base_model.* before loading.
KintsugiBackboneBase		—	—	Base for pretrained backbones: holds self.backbone and adds optional LoRA.
KintsugiBackboneBase.__init__		—	—	No parameters; subclasses build self.backbone.
KintsugiBackboneBase.apply_lora	9	—	BERTBackbone.__init__, LLAMABackbone.__init__, Wav2Vec2Backbone.__init__, WhisperBackbone.__init__	Wrap self.backbone with PEFT LoRA from lora_params; no-op (full fine-tune) if empty.
WhisperBackbone		—	DepAnxClassifier.__init__ mtl_mar_2024/model.py, DepAnxClassifier.__init__ mtl_aug_2024/model.py, DepAnxClassifierWhisperBERT.__init__ mtl_jun_2024/model.py, DepAnxClassifierWhisperLLAMA.__init__ mtl_jun_2024/model.py, WhisperLLMA.__init__ mtl_aug_2024/model.py	Whisper encoder only (decoder dropped), optionally mean-pooled over time.
WhisperBackbone.__init__		KintsugiBackboneBase.apply_lora	—	Load the pretrained Whisper encoder from HuggingFace and apply LoRA.
WhisperBackbone.forward	30	—	—	Whisper encoder on log-mel windows, mean over time -> [ΣW, D] (768 for small.en).
Wav2Vec2Backbone		—	—	wav2vec2 encoder on raw audio, mean-pooled (not used by the stable DAM configs).
Wav2Vec2Backbone.__init__		KintsugiBackboneBase.apply_lora	—	Load the pretrained wav2vec2 model from HuggingFace and apply LoRA.
Wav2Vec2Backbone.forward		—	—	Encode raw audio x["audio"] [ΣW, samples] -> [ΣW, D] (mean over frames if mean_pool).
BERTBackbone		—	DepAnxClassifierWhisperBERT.__init__ mtl_jun_2024/model.py	BERT on transcript tokens; returns the pooler output [ΣW, 768] (DAM 2 text branch).
BERTBackbone.__init__		KintsugiBackboneBase.apply_lora	—	Load pretrained BERT from HuggingFace and apply LoRA.
BERTBackbone.forward		—	—	Pooler output ([CLS] + dense + tanh) of x["text"] tokens. Shapes: -> [ΣW, 768].
LLAMABackbone		—	DepAnxClassifierLLAMA.__init__ mtl_jun_2024/model.py, DepAnxClassifierWhisperLLAMA.__init__ mtl_jun_2024/model.py	LLaMA cut to num_layers layers; embeds text as the last real token's hidden state.
LLAMABackbone.__init__		KintsugiBackboneBase.apply_lora, get_wandb_run_name_tag_from_artifact_url utils.py, localize_wandb_artifact utils.py	—	Load LLaMA (HF id or W&B artifact URL), keep the first num_layers layers, LoRA.
LLAMABackbone.forward		—	—	Last hidden state at the last non-padded token of each sequence.
MultitaskHead		—	DepAnxClassifier.__init__ mtl_aug_2024/model.py, DepAnxClassifier.__init__ mtl_mar_2024/model.py, DepAnxClassifierBase.__init__ mtl_jun_2024/model.py	Shared MLP followed by one small head per task (depression, anxiety).
MultitaskHead.__init__		—	—	Build SharedLayers(backbone_dim, shared_projection_dim) and a TaskHead per task config.
MultitaskHead.forward	31	—	—	Shared layers, then one TaskHead per task -> {task: [ΣW, 1] score}.
ResNetHead		—	—	Alternative head: 2-D ResNet over [N, T, D] features, one logit per task (unused).
ResNetHead.__init__		—	—	ResNet from a HF config, randomly initialized, 1 input channel, len(tasks) outputs.
ResNetHead.forward		—	—	Shapes: [N, T, D] -> {task: [N, 1]}.
ResNetHead1d		—	—	Alternative head: ResNet with feature dim as channels over time (unused).
ResNetHead1d.__init__		—	—	ResNet from a HF config, randomly initialized, backbone_dim channels, len(tasks) outputs.
ResNetHead1d.forward		—	—	Shapes: [N, T, D] -> transposed to [N, D, T, 1] -> {task: [N, 1]}.

kirad/losses.py

CORAL ordinal-regression loss (stage: loss). · open file ↗

Function	Step	Calls	Called by	What it does
CORALLoss		—	OrdinalRegressionPLModule.__init__ base_models.py	CORAL loss: score + learned per-cutoff biases, BCE against the thermometer code of y.
CORALLoss.__init__		—	—	One learned bias per cutoff (num_classes - 1, e.g. 27 for PHQ-9) and a 0/1 weight each.
CORALLoss.forward	35	—	—	Weighted sum of BCE over cutoffs: sigmoid(s + b_k) vs [y > k] for each cutoff k.

kirad/ordinal_thresholding.py

Tune ordinal thresholds that turn continuous scores into severity classes. · open file ↗

Function	Step	Calls	Called by	What it does
OrdinalThresholding		—	—	Base module: K-1 sorted thresholds (buffer) that map scores to K ordinal classes.
OrdinalThresholding.__init__		—	—	Create num_classes - 1 zero thresholds as a buffer (saved in the checkpoint).
OrdinalThresholding.is_valid		—	OrdinalRegressionPLModule.log_thresholding_metrics base_models.py	Check whether the thresholds are monotone non-decreasing.
OrdinalThresholding.forward		—	—	Class index per score = number of thresholds strictly below it (searchsorted).
OrdinalThresholding.tune_thresholds		—	—	Tune thresholds on validation scores/labels; no-op here, overridden by subclasses.
CORALThresholding		—	OrdinalRegressionPLModule.__init__ base_models.py	Thresholds = -CORAL bias at the chosen cutoff indices (learned on train, no tuning).
CORALThresholding.__init__		—	—	Keep a reference to the CORAL biases and the cutoff indices to read.
CORALThresholding.tune_thresholds		—	—	Copy -biases[threshold_indices] from the CORAL loss; scores and labels are ignored.
OptimalOrdinalThresholdingViaDynamicProgramming		—	—	Base for DP threshold tuning: pick K-1 thresholds from candidates to optimize a metric.
OptimalOrdinalThresholdingViaDynamicProgramming.__init__		—	—	Check that the subclass set direction to "min" or "max".
OptimalOrdinalThresholdingViaDynamicProgramming.mean_cost		—	OptimalOrdinalThresholdingViaDynamicProgramming.best_constant_output_classifier	Compute the mean cost of assigning label(s) preds when the ground truth is labels.
OptimalOrdinalThresholdingViaDynamicProgramming.best_constant_output_classifier		OptimalOrdinalThresholdingViaDynamicProgramming.mean_cost	OrdinalRegressionPLModule.log_thresholding_metrics base_models.py	Best metric of a constant-class predictor and that class (baseline_const_* in logs).
OptimalOrdinalThresholdingViaDynamicProgramming.dp_step		—	OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds	One DP step: best cost of classes <= c_idx for each candidate upper threshold of c_idx.
OptimalOrdinalThresholdingViaDynamicProgramming.tune_thresholds	40	OptimalOrdinalThresholdingViaDynamicProgramming.dp_step	—	Set thresholds by DP over midpoints of unique scores (+/-inf); return the best metric.
cumsum_with_0		—	OptimalCostPerSampleOrdinalThresholding.dp_step, _compute_metrics_matrices	Cumulative sum along dim 0 with a leading 0 (length n + 1).
OptimalCostPerSampleOrdinalThresholding		—	—	DP tuning for metrics that are a mean of per-sample costs (accuracy, MAE, macro recall).
OptimalCostPerSampleOrdinalThresholding.cost		—	OptimalCostPerSampleOrdinalThresholding.dp_step, OptimalCostPerSampleOrdinalThresholding.mean_cost	Per-sample cost of predicting preds for labels (subclasses define the metric).
OptimalCostPerSampleOrdinalThresholding.mean_cost		OptimalCostPerSampleOrdinalThresholding.cost	—	Compute the mean cost of assigning label(s) preds when the ground truth is labels.
OptimalCostPerSampleOrdinalThresholding.dp_step		OptimalCostPerSampleOrdinalThresholding.cost, cumsum_with_0	—	O(n) DP step: histogram per-sample costs over candidates, cumsum, running max.
MaxAccuracyOrdinalThresholding		—	OrdinalRegressionPLModule.__init__ base_models.py	Threshold to maximize accuracy.
MaxAccuracyOrdinalThresholding.cost		—	—	1.0 where pred == label, else 0.0.
MaxMacroRecallOrdinalThresholding		—	OrdinalRegressionPLModule.__init__ base_models.py	Threshold to maximize macro-averaged recall.
MaxMacroRecallOrdinalThresholding.cost		—	—	Correct = N / (K * count[label]), so the mean equals macro-averaged recall.
MinAbsoluteErrorOrdinalThresholding		—	OrdinalRegressionPLModule.__init__ base_models.py	Threshold to minimize mean absolute error.
MinAbsoluteErrorOrdinalThresholding.cost		—	—	|pred - label| in class units.
ClassWeightedOptimalCostPerSampleOrdinalThresholding		—	—	Wrap a per-sample cost so each class weighs equally; not used by base_models.
ClassWeightedOptimalCostPerSampleOrdinalThresholding.__init__		—	—	Copy direction and num_classes from the wrapped instance.
ClassWeightedOptimalCostPerSampleOrdinalThresholding.cost		—	—	Wrapped cost times N / (K * count[label]); raises if any class is missing.
OptimalCostPerClassOrdinalThresholding		—	—	DP tuning for per-class metrics (macro precision, macro F1) via a T x T cost matrix.
OptimalCostPerClassOrdinalThresholding.cost_matrix		—	OptimalCostPerClassOrdinalThresholding.dp_step, OptimalCostPerClassOrdinalThresholding.mean_cost	Cost matrix [T, T]: metric of class c_idx if scores in [thr[i], thr[j]) get c_idx.
OptimalCostPerClassOrdinalThresholding.mean_cost		OptimalCostPerClassOrdinalThresholding.cost_matrix	—	Macro metric of fixed preds, reusing cost_matrix with thresholds c +/- 0.5.
OptimalCostPerClassOrdinalThresholding.dp_step		OptimalCostPerClassOrdinalThresholding.cost_matrix	—	O(T^2) DP step: add prev_cost to each lower-threshold row and max over it.
_compute_metrics_matrices		cumsum_with_0	MaxMacroF1OrdinalThresholding.cost_matrix, MaxMacroPrecisionOrdinalThresholding.cost_matrix	TP and predicted-positive counts for every [thresholds[i], thresholds[j]) range.
MaxMacroPrecisionOrdinalThresholding		—	OrdinalRegressionPLModule.__init__ base_models.py	Threshold to maximize macro-averaged precision.
MaxMacroPrecisionOrdinalThresholding.cost_matrix		_compute_metrics_matrices	—	Precision of class c_idx for each (lower, upper) threshold pair; empty range = 0.
MaxMacroF1OrdinalThresholding		—	OrdinalRegressionPLModule.__init__ base_models.py	Threshold to maximize macro-averaged F1 score.
MaxMacroF1OrdinalThresholding.cost_matrix		_compute_metrics_matrices	—	F1 of class c_idx = 2TP / (TP+FP + TP+FN) for each (lower, upper) threshold pair.

kirad/metrics.py

Binary Sn/Sp metrics with an indeterminate band, plus a DDP-safe score accumulator. · open file ↗

Function	Step	Calls	Called by	What it does
pad_cat		—	—	Zero-pad non-batch dims to the largest shape, then concat on dim 0 (DDP reduce).
MultiCatMetric		—	OrdinalRegressionPLModule.__init__ base_models.py	Accumulate parallel tensors over batches and ranks (uuids, scores, variances, labels).
MultiCatMetric.__init__		—	—	Register num_states list states, reduced across ranks with pad_cat.
MultiCatMetric.update		—	—	Append one batch of num_states tensors.
MultiCatMetric.compute		—	—	Concatenate each of the num_states tensors over all calls to update.
running_argmax_indices		—	pareto_2d_indices	Return indices of a where the value is larger than all previous values.
pareto_2d_indices		running_argmax_indices	IndetSnSpArray.roc_curve, IndetSnSpArray.sn_eq_sp_graph	Indices of the 2-d Pareto frontier maximizing x and y.
midpoints_with_infs		—	IndetRocCurve.build	Return the midpoints between the sorted unique elements of x, along with +/-inf.
IndetSnSpArray		—	IndetSnSpArray.__getitem__, IndetSnSpTensorModule.numpy	Sn, Sp and indeterminate fraction for many (lower, upper) threshold pairs.
IndetSnSpArray.min_sn_sp		—	—	min(Sn, Sp) per threshold pair; the quantity maximized for sn_eq_sp.
IndetSnSpArray.build	39	—	—	Sn/Sp/indet_frac for all threshold pairs (or the given ones) via histogram cumsums.
IndetSnSpArray.eval		—	—	Evaluate the given data on the thresholds of self (e.g. val thresholds on test).
IndetSnSpArray.__getitem__		IndetSnSpArray	—	Extract a subarray with numpy-style indexing.
IndetSnSpArray.roc_curve		IndetRocCurve, pareto_2d_indices	OrdinalRegressionPLModule.log_ordinal_metrics base_models.py, indet_analysis utils.py, plot_indeterminate_analysis plotting.py	Pareto ROC curve of (Sn, Sp) using only pairs with indet_frac <= indet_budget.
IndetSnSpArray.sn_eq_sp_graph		IndetSnEqSpGraph, pareto_2d_indices	OrdinalRegressionPLModule.on_validation_test_epoch_end base_models.py, plot_indeterminate_analysis plotting.py	Best min(Sn, Sp) per indeterminate fraction; these thresholds are saved in the ckpt.
IndetSnSpTensorModule		—	OrdinalRegressionPLModule.__init__ base_models.py	Buffers holding an IndetSnSpArray (val sn_eq_sp thresholds) in the checkpoint.
IndetSnSpTensorModule.__init__		—	—	One empty buffer per IndetSnSpArray field.
IndetSnSpTensorModule.update		—	—	Update the stored values to those given in the array.
IndetSnSpTensorModule.numpy		IndetSnSpArray	—	Return an IndetSnSpArray with values taken from the stored tensors.
IndetRocCurve		—	IndetSnSpArray.roc_curve	ROC curve (Pareto Sn/Sp points) within an indeterminate budget, with thresholds.
IndetRocCurve.sn_eq_sp		—	OrdinalRegressionPLModule.log_ordinal_metrics base_models.py, analyze_dataframe_with_scores utils.py, indet_analysis utils.py	Locate the point on the ROC curve closest to the diagonal (max min(Sn, Sp)).
IndetRocCurve.auc		—	OrdinalRegressionPLModule.log_ordinal_metrics base_models.py, analyze_dataframe_with_scores utils.py, indet_analysis utils.py	Compute the area under the ROC curve.
IndetRocCurve.build		midpoints_with_infs	—	Plain ROC curve (no indeterminate band) in O(n log n).
IndetSnEqSpGraph		—	IndetSnSpArray.sn_eq_sp_graph, OrdinalRegressionPLModule.log_ordinal_metrics base_models.py	Best Sn=Sp for each indeterminate budget, with its lower/upper thresholds.
IndetSnEqSpGraph.at_budget		—	OrdinalRegressionPLModule.log_ordinal_metrics base_models.py	Locate the best point on the graph within the given budget.
binary_indet_confusion_matrix		—	—	Flat 2x3 confusion matrix (neg/pos truth x neg/pos/indet output) for one threshold pair.
bootstrap_confusion_matrix		—	—	Multinomial bootstrap of a flat confusion matrix; returns len(cm) x num_bootstrap.

kirad/plotting.py

Matplotlib plots of the indeterminate-band trade-off for one task's scores. · open file ↗

Function	Step	Calls	Called by	What it does
plot_indeterminate_analysis		IndetSnSpArray.roc_curve metrics.py, IndetSnSpArray.sn_eq_sp_graph metrics.py	—	Plot thresholds and Sn=Sp vs indeterminate budget, plus ROC curves per binary cut.

research/stable/ordinal_regression_mtl_jun_2024/model.py

DAM 2 models (stage: model): audio + transcript text classifiers. · open file ↗

Function	Step	Calls	Called by	What it does
create_task_modules_from_params		OrdinalRegressionPLModule base_models.py	DepAnxClassifierBase.__init__	One OrdinalRegressionPLModule per (task, config); cutoffs from kirad/constants.py.
DepAnxClassifierBase		—	—	Concatenate embeddings of several backbones (ModuleDict order) and feed a MultitaskHead.
DepAnxClassifierBase.__init__		MultitaskHead base_models.py, create_task_modules_from_params	—	Build per-task modules and a MultitaskHead over the summed backbone dims.
DepAnxClassifierBase.forward		DepAnxClassifierBase.forward_backbone	—	Backbones then MultitaskHead -> ({task: [ΣW, 1] scores}, lengths).
DepAnxClassifierBase.forward_backbone		—	DepAnxClassifierBase.compute_features_to_cache, DepAnxClassifierBase.forward	Run each backbone on the batch and concatenate on dim 1 -> [ΣW, ΣD].
DepAnxClassifierBase.compute_features_to_cache		DepAnxClassifierBase.forward_backbone	—	Concatenated embeddings split per recording (tuple of [W_i, ΣD]) for cache.
DepAnxClassifierBase.configure_optimizers		—	—	AdamW with two groups (backbone, head + CORAL biases) and MultiStepLR.
DepAnxClassifierWhisperBERT		—	—	DAM 2: Whisper audio + BERT transcript embeddings (768 + 768) -> multitask head.
DepAnxClassifierWhisperBERT.__init__		BERTBackbone base_models.py, WhisperBackbone base_models.py	—	Backbones {"audio": WhisperBackbone, "text": BERTBackbone}; order sets the concat order.
DepAnxClassifierWhisperLLAMA		—	—	Whisper audio + LLaMA transcript embeddings concatenated -> multitask head.
DepAnxClassifierWhisperLLAMA.__init__		LLAMABackbone base_models.py, WhisperBackbone base_models.py	—	Backbones {"audio": WhisperBackbone, "text": LLAMABackbone}; order sets the concat order.
DepAnxClassifierLLAMA		—	—	Text only: LLaMA transcript embedding -> multitask head.
DepAnxClassifierLLAMA.__init__		LLAMABackbone base_models.py	—	Single backbone {"text": LLAMABackbone}.

research/stable/ordinal_regression_mtl_aug_2024/model.py

DAM 3 models (stage: model): audio-only model that imitates DAM 2's text branch. · open file ↗

Function	Step	Calls	Called by	What it does
WhisperLLMA		—	DepAnxClassifier.__init__	DAM 3 phase 1: Whisper encoder trained to predict a cached text embedding from audio.
WhisperLLMA.__init__		WhisperBackbone base_models.py	—	Whisper backbone "llma" plus a fixed random projection of the target if dims differ.
WhisperLLMA.forward		average_tensor_in_segments utils.py	WhisperLLMA.step	Whisper embedding averaged per recording. Shapes: log-mel [ΣW, 80, 3000] -> [B, 768].
WhisperLLMA.forward_ground_truth		average_tensor_in_segments utils.py	WhisperLLMA.step	Target: cached text embedding (x["backbone_cache"]) averaged per recording, projected.
WhisperLLMA.training_step		WhisperLLMA.step	—	Lightning train step: returns the MSE loss.
WhisperLLMA.validation_step		WhisperLLMA.step	—	Lightning val step; loss is logged only.
WhisperLLMA.test_step		WhisperLLMA.step	—	Lightning test step; loss is logged only.
WhisperLLMA.step		WhisperLLMA.forward, WhisperLLMA.forward_ground_truth	WhisperLLMA.test_step, WhisperLLMA.training_step, WhisperLLMA.validation_step	MSE between per-recording Whisper embedding and the cached teacher embedding.
WhisperLLMA.configure_optimizers		—	—	AdamW on the Whisper backbone only (optimizer["backbone"]) with MultiStepLR.
DepAnxClassifier		—	—	DAM 3: frozen Whisper + frozen WhisperLLMA embeddings (768 + 768) -> trainable head.
DepAnxClassifier.__init__		MultitaskHead base_models.py, OrdinalRegressionPLModule base_models.py, WhisperBackbone base_models.py, WhisperLLMA, get_wandb_run_name_tag_from_artifact_url utils.py, localize_wandb_artifact utils.py	—	Build per-task modules, a Whisper audio backbone, load WhisperLLMA from W&B, head.
DepAnxClassifier.forward		DepAnxClassifier.forward_backbone	—	Frozen backbones then MultitaskHead -> ({task: [B, 1] scores}, lengths of 1).
DepAnxClassifier.forward_backbone		average_tensor_in_segments utils.py	DepAnxClassifier.compute_features_to_cache, DepAnxClassifier.forward	No-grad Whisper + WhisperLLMA, averaged per recording -> [B, 1536]; lengths set to 1.
DepAnxClassifier.compute_features_to_cache		DepAnxClassifier.forward_backbone	—	Per-recording embeddings (one [1, 1536] each) for cache.
DepAnxClassifier.configure_optimizers		—	—	AdamW on head + CORAL biases only (backbones frozen), with MultiStepLR.