src/bin/prepare_external_datasets.rs

#![deny(missing_docs)]
//! Download public optional DADOES training datasets.

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

const DEFAULT_EXTERNAL_DATA_DIR: &str = "data/raw/external";

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

fn run() -> Result<(), PrepareDatasetsError> {
    let args = PrepareDatasetsArgs::from_env()?;
    for source in PUBLIC_DOWNLOADS {
        let target = args.root.join(source.relative_path);
        download_file(source.url, &target)?;
        println!("wrote {}", target.display());
    }

    Ok(())
}

fn download_file(url: &'static str, target: &Path) -> Result<(), PrepareDatasetsError> {
    create_parent_dir(target)?;
    let response = ureq::get(url)
        .call()
        .map_err(|source| PrepareDatasetsError::Download {
            url,
            source: Box::new(source),
        })?;
    let mut reader = response.into_reader();
    let mut output = File::create(target).map_err(|source| PrepareDatasetsError::WriteFile {
        path: target.to_path_buf(),
        source,
    })?;
    std::io::copy(&mut reader, &mut output).map_err(|source| PrepareDatasetsError::WriteFile {
        path: target.to_path_buf(),
        source,
    })?;
    Ok(())
}

fn create_parent_dir(path: &Path) -> Result<(), PrepareDatasetsError> {
    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| PrepareDatasetsError::WriteFile {
        path: parent.to_path_buf(),
        source,
    })
}

#[derive(Debug, Clone, Eq, PartialEq)]
struct PrepareDatasetsArgs {
    root: PathBuf,
}

impl PrepareDatasetsArgs {
    fn from_env() -> Result<Self, PrepareDatasetsError> {
        let args = std::env::args().skip(1).collect::<Vec<_>>();
        match args.as_slice() {
            [] => Ok(Self {
                root: PathBuf::from(DEFAULT_EXTERNAL_DATA_DIR),
            }),
            [root] => Ok(Self {
                root: PathBuf::from(root),
            }),
            _ => Err(PrepareDatasetsError::Usage),
        }
    }
}

#[derive(Debug)]
enum PrepareDatasetsError {
    Usage,
    Download {
        url: &'static str,
        source: Box<ureq::Error>,
    },
    WriteFile {
        path: PathBuf,
        source: std::io::Error,
    },
}

impl Display for PrepareDatasetsError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            PrepareDatasetsError::Usage => write!(
                formatter,
                "usage: cargo run --release --bin prepare_external_datasets -- [output_dir]"
            ),
            PrepareDatasetsError::Download { url, source } => {
                write!(formatter, "failed to download {url}: {source}")
            }
            PrepareDatasetsError::WriteFile { path, source } => {
                write!(formatter, "failed to write {}: {source}", path.display())
            }
        }
    }
}

impl Error for PrepareDatasetsError {}

#[derive(Debug, Copy, Clone)]
struct PublicDownload {
    url: &'static str,
    relative_path: &'static str,
}

const PUBLIC_DOWNLOADS: &[PublicDownload] = &[
    PublicDownload {
        url: "https://huggingface.co/datasets/yael-katsman/Loneliness-Causes-and-Intensity/resolve/main/train_data.csv",
        relative_path: "loneliness_causes/train_data.csv",
    },
    PublicDownload {
        url: "https://huggingface.co/datasets/yael-katsman/Loneliness-Causes-and-Intensity/resolve/main/validation_data.csv",
        relative_path: "loneliness_causes/validation_data.csv",
    },
    PublicDownload {
        url: "https://huggingface.co/datasets/yael-katsman/Loneliness-Causes-and-Intensity/resolve/main/test_data.csv",
        relative_path: "loneliness_causes/test_data.csv",
    },
    PublicDownload {
        url: "https://huggingface.co/datasets/smit18/text_emotion/resolve/main/text_emotion.csv",
        relative_path: "text_emotion/text_emotion.csv",
    },
    PublicDownload {
        url: "https://huggingface.co/datasets/PrajwalNayaka/Text-Emotion/resolve/main/final_dataset.csv",
        relative_path: "prajwal_text_emotion/final_dataset.csv",
    },
    PublicDownload {
        url: "https://huggingface.co/datasets/Um1neko/text_emotion/resolve/main/train.json",
        relative_path: "um1neko/train.json",
    },
    PublicDownload {
        url: "https://huggingface.co/datasets/Um1neko/text_emotion/resolve/main/dev.json",
        relative_path: "um1neko/dev.json",
    },
    PublicDownload {
        url: "https://huggingface.co/datasets/Um1neko/text_emotion/resolve/main/test.json",
        relative_path: "um1neko/test.json",
    },
];