src/bin/train.rs

#![deny(missing_docs)]
//! Train the DADOES linear mood classifier from GoEmotions splits.

use std::error::Error;
use std::fmt::{Display, Formatter};
use std::fs::{self, File};
use std::path::{Path, PathBuf};
use std::process::ExitCode;

use dadoes::datasets::{
    DatasetParseError, parse_dadoes_jsonl, parse_empathetic_dialogues_csv,
    parse_fig_loneliness_jsonl, parse_loneliness_causes_csv, parse_prajwal_text_emotion_csv,
    parse_text_emotion_csv, parse_um1neko_text_emotion_jsonl,
};
use dadoes::goemotions::{GoEmotionsError, parse_split};
use dadoes::{
    EvaluationMetrics, LinearMoodModel, ModelError, ModelIoError, OwnedTrainingExample,
    TrainConfig, owned_seed_training_examples,
};

const DEFAULT_DATA_DIR: &str = "data/raw/goemotions";
const DEFAULT_EXTERNAL_DATA_DIR: &str = "data/raw/external";
const DEFAULT_MODEL_PATH: &str = "models/goemotions-linear.dadoes";
const DEFAULT_REPORT_PATH: &str = "models/goemotions-linear.report.json";
const EXTERNAL_DATA_DIR_ENV: &str = "DADOES_EXTERNAL_DATA_DIR";
const INCLUDE_NON_COMMERCIAL_ENV: &str = "DADOES_INCLUDE_NON_COMMERCIAL";
const SEED_AUGMENTATION_REPEATS: usize = 4;

/// Runs the DADOES training command.
fn main() -> ExitCode {
    match run() {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::from(1)
        }
    }
}

fn run() -> Result<(), TrainCliError> {
    let args = TrainArgs::from_env()?;
    let mut train = load_split(&args.data_dir, "train.tsv")?;
    let dev = load_split(&args.data_dir, "dev.tsv")?;
    let test = load_split(&args.data_dir, "test.tsv")?;
    let mut dev = dev;
    let mut test = test;
    let external = load_external_dataset_mix(
        &args.external_data_dir,
        args.include_non_commercial,
        &mut train,
        &mut dev,
        &mut test,
    )?;
    let seed_examples = owned_seed_training_examples()?;
    let seed_count = seed_examples.len();
    for _ in 0..SEED_AUGMENTATION_REPEATS {
        train.extend(seed_examples.iter().cloned());
    }
    let config = TrainConfig::default();

    println!(
        "loaded splits: train={} dev={} test={} seed_examples={} seed_repeats={}",
        train.len(),
        dev.len(),
        test.len(),
        seed_count,
        SEED_AUGMENTATION_REPEATS
    );
    println!(
        "external_data_dir={} include_non_commercial={}",
        args.external_data_dir.display(),
        args.include_non_commercial
    );
    if external.is_empty() {
        println!("external datasets: no optional files found");
    } else {
        println!("external datasets:");
        for summary in &external {
            println!(
                "- {} {} examples={} path={}",
                summary.name,
                summary.split.as_str(),
                summary.examples,
                summary.path.display()
            );
        }
    }
    println!(
        "training: features={} epochs={} lr={} l2={} patience={} threshold={}",
        config.feature_count,
        config.epochs,
        config.learning_rate,
        config.l2,
        config.patience,
        config.threshold
    );

    let outcome = LinearMoodModel::train_with_validation(&train, &dev, config)?;
    let test_metrics = outcome.model.evaluate(&test, config.threshold)?;
    create_parent_dir(&args.model_path)?;
    let mut model_file =
        File::create(&args.model_path).map_err(|source| TrainCliError::WriteFile {
            path: args.model_path.clone(),
            source,
        })?;
    outcome.model.save_to_writer(&mut model_file)?;

    let report = report_json(
        &outcome.validation,
        &test_metrics,
        outcome.epochs_trained,
        outcome.best_epoch,
        outcome.best_validation_loss,
    );
    create_parent_dir(&args.report_path)?;
    fs::write(&args.report_path, report).map_err(|source| TrainCliError::WriteFile {
        path: args.report_path.clone(),
        source,
    })?;

    println!(
        "best_epoch={} epochs_trained={} dev_loss={:.6} dev_micro_f1={:.4}",
        outcome.best_epoch,
        outcome.epochs_trained,
        outcome.validation.loss,
        outcome.validation.micro_f1
    );
    println!(
        "test_loss={:.6} test_micro_f1={:.4} test_exact_match={:.4}",
        test_metrics.loss, test_metrics.micro_f1, test_metrics.exact_match
    );
    println!("saved model: {}", args.model_path.display());
    println!("saved report: {}", args.report_path.display());

    Ok(())
}

fn load_external_dataset_mix(
    root: &Path,
    include_non_commercial: bool,
    train: &mut Vec<OwnedTrainingExample>,
    dev: &mut Vec<OwnedTrainingExample>,
    test: &mut Vec<OwnedTrainingExample>,
) -> Result<Vec<ExternalDatasetSummary>, TrainCliError> {
    let mut summaries = Vec::new();
    for file in EXTERNAL_DATASET_FILES {
        if !file.license.allowed_by(include_non_commercial) {
            continue;
        }

        let path = root.join(file.relative_path);
        if !path.exists() {
            continue;
        }

        let input = fs::read_to_string(&path).map_err(|source| TrainCliError::ReadFile {
            path: path.clone(),
            source,
        })?;
        let examples = (file.parser)(&input).map_err(|source| TrainCliError::ParseExternal {
            path: path.clone(),
            source,
        })?;
        let examples_count = examples.len();
        match file.split {
            DatasetSplit::Train => train.extend(examples),
            DatasetSplit::Dev => dev.extend(examples),
            DatasetSplit::Test => test.extend(examples),
        }
        summaries.push(ExternalDatasetSummary {
            name: file.name,
            split: file.split,
            path,
            examples: examples_count,
        });
    }

    Ok(summaries)
}

fn load_split(
    data_dir: &Path,
    file_name: &'static str,
) -> Result<Vec<OwnedTrainingExample>, TrainCliError> {
    let path = data_dir.join(file_name);
    let input = fs::read_to_string(&path).map_err(|source| TrainCliError::ReadFile {
        path: path.clone(),
        source,
    })?;
    parse_split(&input).map_err(|source| TrainCliError::ParseSplit {
        split: file_name,
        source,
    })
}

fn create_parent_dir(path: &Path) -> Result<(), TrainCliError> {
    let Some(parent) = path.parent() else {
        return Ok(());
    };
    if parent.as_os_str().is_empty() {
        return Ok(());
    }

    fs::create_dir_all(parent).map_err(|source| TrainCliError::WriteFile {
        path: parent.to_path_buf(),
        source,
    })
}

fn report_json(
    validation: &EvaluationMetrics,
    test: &EvaluationMetrics,
    epochs_trained: usize,
    best_epoch: usize,
    best_validation_loss: f32,
) -> String {
    format!(
        concat!(
            "{{\n",
            "  \"epochs_trained\": {},\n",
            "  \"best_epoch\": {},\n",
            "  \"best_validation_loss\": {:.6},\n",
            "  \"validation\": {},\n",
            "  \"test\": {}\n",
            "}}\n"
        ),
        epochs_trained,
        best_epoch,
        best_validation_loss,
        metrics_json(validation),
        metrics_json(test)
    )
}

fn metrics_json(metrics: &EvaluationMetrics) -> String {
    format!(
        concat!(
            "{{",
            "\"examples\":{},",
            "\"loss\":{:.6},",
            "\"micro_precision\":{:.6},",
            "\"micro_recall\":{:.6},",
            "\"micro_f1\":{:.6},",
            "\"exact_match\":{:.6}",
            "}}"
        ),
        metrics.examples,
        metrics.loss,
        metrics.micro_precision,
        metrics.micro_recall,
        metrics.micro_f1,
        metrics.exact_match
    )
}

#[derive(Debug, Clone, Eq, PartialEq)]
struct TrainArgs {
    data_dir: PathBuf,
    external_data_dir: PathBuf,
    include_non_commercial: bool,
    model_path: PathBuf,
    report_path: PathBuf,
}

impl TrainArgs {
    fn from_env() -> Result<Self, TrainCliError> {
        let args: Vec<String> = std::env::args().skip(1).collect();
        let external_data_dir = std::env::var(EXTERNAL_DATA_DIR_ENV)
            .map(PathBuf::from)
            .unwrap_or_else(|_| PathBuf::from(DEFAULT_EXTERNAL_DATA_DIR));
        let include_non_commercial = env_flag(INCLUDE_NON_COMMERCIAL_ENV);
        match args.as_slice() {
            [] => Ok(Self {
                data_dir: PathBuf::from(DEFAULT_DATA_DIR),
                external_data_dir,
                include_non_commercial,
                model_path: PathBuf::from(DEFAULT_MODEL_PATH),
                report_path: PathBuf::from(DEFAULT_REPORT_PATH),
            }),
            [data_dir, model_path] => Ok(Self {
                data_dir: PathBuf::from(data_dir),
                external_data_dir,
                include_non_commercial,
                model_path: PathBuf::from(model_path),
                report_path: PathBuf::from(DEFAULT_REPORT_PATH),
            }),
            [data_dir, model_path, report_path] => Ok(Self {
                data_dir: PathBuf::from(data_dir),
                external_data_dir,
                include_non_commercial,
                model_path: PathBuf::from(model_path),
                report_path: PathBuf::from(report_path),
            }),
            _ => Err(TrainCliError::Usage),
        }
    }
}

fn env_flag(name: &str) -> bool {
    match std::env::var(name) {
        Ok(value) => matches!(
            value.trim().to_ascii_lowercase().as_str(),
            "1" | "true" | "yes" | "on"
        ),
        Err(_) => false,
    }
}

type ExternalParser = fn(&str) -> Result<Vec<OwnedTrainingExample>, DatasetParseError>;

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum DatasetSplit {
    Train,
    Dev,
    Test,
}

impl DatasetSplit {
    fn as_str(self) -> &'static str {
        match self {
            DatasetSplit::Train => "train",
            DatasetSplit::Dev => "dev",
            DatasetSplit::Test => "test",
        }
    }
}

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum LicenseClass {
    Public,
    NonCommercial,
}

impl LicenseClass {
    fn allowed_by(self, include_non_commercial: bool) -> bool {
        match self {
            LicenseClass::Public => true,
            LicenseClass::NonCommercial => include_non_commercial,
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct ExternalDatasetFile {
    name: &'static str,
    license: LicenseClass,
    split: DatasetSplit,
    relative_path: &'static str,
    parser: ExternalParser,
}

#[derive(Debug, Clone)]
struct ExternalDatasetSummary {
    name: &'static str,
    split: DatasetSplit,
    path: PathBuf,
    examples: usize,
}

const EXTERNAL_DATASET_FILES: &[ExternalDatasetFile] = &[
    ExternalDatasetFile {
        name: "dadoes-domain",
        license: LicenseClass::Public,
        split: DatasetSplit::Train,
        relative_path: "dadoes-domain/train.jsonl",
        parser: parse_dadoes_jsonl,
    },
    ExternalDatasetFile {
        name: "dadoes-domain",
        license: LicenseClass::Public,
        split: DatasetSplit::Dev,
        relative_path: "dadoes-domain/dev.jsonl",
        parser: parse_dadoes_jsonl,
    },
    ExternalDatasetFile {
        name: "dadoes-domain",
        license: LicenseClass::Public,
        split: DatasetSplit::Test,
        relative_path: "dadoes-domain/test.jsonl",
        parser: parse_dadoes_jsonl,
    },
    ExternalDatasetFile {
        name: "facebook/empathetic_dialogues",
        license: LicenseClass::NonCommercial,
        split: DatasetSplit::Train,
        relative_path: "empathetic_dialogues/train.csv",
        parser: parse_empathetic_dialogues_csv,
    },
    ExternalDatasetFile {
        name: "facebook/empathetic_dialogues",
        license: LicenseClass::NonCommercial,
        split: DatasetSplit::Dev,
        relative_path: "empathetic_dialogues/validation.csv",
        parser: parse_empathetic_dialogues_csv,
    },
    ExternalDatasetFile {
        name: "facebook/empathetic_dialogues",
        license: LicenseClass::NonCommercial,
        split: DatasetSplit::Test,
        relative_path: "empathetic_dialogues/test.csv",
        parser: parse_empathetic_dialogues_csv,
    },
    ExternalDatasetFile {
        name: "FIG-Loneliness/FIG-Loneliness",
        license: LicenseClass::NonCommercial,
        split: DatasetSplit::Train,
        relative_path: "fig_loneliness/train.jsonl",
        parser: parse_fig_loneliness_jsonl,
    },
    ExternalDatasetFile {
        name: "FIG-Loneliness/FIG-Loneliness",
        license: LicenseClass::NonCommercial,
        split: DatasetSplit::Dev,
        relative_path: "fig_loneliness/dev.jsonl",
        parser: parse_fig_loneliness_jsonl,
    },
    ExternalDatasetFile {
        name: "FIG-Loneliness/FIG-Loneliness",
        license: LicenseClass::NonCommercial,
        split: DatasetSplit::Test,
        relative_path: "fig_loneliness/test.jsonl",
        parser: parse_fig_loneliness_jsonl,
    },
    ExternalDatasetFile {
        name: "yael-katsman/Loneliness-Causes-and-Intensity",
        license: LicenseClass::Public,
        split: DatasetSplit::Train,
        relative_path: "loneliness_causes/train_data.csv",
        parser: parse_loneliness_causes_csv,
    },
    ExternalDatasetFile {
        name: "yael-katsman/Loneliness-Causes-and-Intensity",
        license: LicenseClass::Public,
        split: DatasetSplit::Dev,
        relative_path: "loneliness_causes/validation_data.csv",
        parser: parse_loneliness_causes_csv,
    },
    ExternalDatasetFile {
        name: "yael-katsman/Loneliness-Causes-and-Intensity",
        license: LicenseClass::Public,
        split: DatasetSplit::Test,
        relative_path: "loneliness_causes/test_data.csv",
        parser: parse_loneliness_causes_csv,
    },
    ExternalDatasetFile {
        name: "smit18/text_emotion",
        license: LicenseClass::Public,
        split: DatasetSplit::Train,
        relative_path: "text_emotion/text_emotion.csv",
        parser: parse_text_emotion_csv,
    },
    ExternalDatasetFile {
        name: "PrajwalNayaka/Text-Emotion",
        license: LicenseClass::Public,
        split: DatasetSplit::Train,
        relative_path: "prajwal_text_emotion/final_dataset.csv",
        parser: parse_prajwal_text_emotion_csv,
    },
    ExternalDatasetFile {
        name: "Um1neko/text_emotion",
        license: LicenseClass::Public,
        split: DatasetSplit::Train,
        relative_path: "um1neko/train.json",
        parser: parse_um1neko_text_emotion_jsonl,
    },
    ExternalDatasetFile {
        name: "Um1neko/text_emotion",
        license: LicenseClass::Public,
        split: DatasetSplit::Dev,
        relative_path: "um1neko/dev.json",
        parser: parse_um1neko_text_emotion_jsonl,
    },
    ExternalDatasetFile {
        name: "Um1neko/text_emotion",
        license: LicenseClass::Public,
        split: DatasetSplit::Test,
        relative_path: "um1neko/test.json",
        parser: parse_um1neko_text_emotion_jsonl,
    },
];

#[derive(Debug)]
enum TrainCliError {
    Usage,
    ReadFile {
        path: PathBuf,
        source: std::io::Error,
    },
    WriteFile {
        path: PathBuf,
        source: std::io::Error,
    },
    ParseSplit {
        split: &'static str,
        source: GoEmotionsError,
    },
    ParseExternal {
        path: PathBuf,
        source: DatasetParseError,
    },
    Model(ModelError),
    Checkpoint(ModelIoError),
}

impl Display for TrainCliError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            TrainCliError::Usage => write!(
                formatter,
                "usage: cargo run --release --bin train -- [data_dir model_path [report_path]]"
            ),
            TrainCliError::ReadFile { path, source } => {
                write!(formatter, "failed to read {}: {source}", path.display())
            }
            TrainCliError::WriteFile { path, source } => {
                write!(formatter, "failed to write {}: {source}", path.display())
            }
            TrainCliError::ParseSplit { split, source } => {
                write!(formatter, "failed to parse {split}: {source}")
            }
            TrainCliError::ParseExternal { path, source } => {
                write!(
                    formatter,
                    "failed to parse external dataset {}: {source}",
                    path.display()
                )
            }
            TrainCliError::Model(error) => write!(formatter, "training failed: {error}"),
            TrainCliError::Checkpoint(error) => write!(formatter, "checkpoint failed: {error}"),
        }
    }
}

impl Error for TrainCliError {}

impl From<ModelError> for TrainCliError {
    fn from(error: ModelError) -> Self {
        TrainCliError::Model(error)
    }
}

impl From<ModelIoError> for TrainCliError {
    fn from(error: ModelIoError) -> Self {
        TrainCliError::Checkpoint(error)
    }
}