use csv::ReaderBuilder;
use dataset_core::{Dataset, DatasetError, acquire_dataset, download_to, unzip};
use ndarray::Array1;
use std::fs::File;
use std::io::Write as _;
type YoutubeSpamData = (Array1<String>, Array1<&'static str>);
const YOUTUBE_SPAM_DATA_URL: &str = "https://archive.ics.uci.edu/ml/machine-learning-databases/00380/YouTube-Spam-Collection-v1.zip";
const YOUTUBE_SPAM_ZIP_FILENAME: &str = "YouTube-Spam-Collection-v1.zip";
const YOUTUBE_SPAM_SOURCE_FILENAMES: [&str; 5] = [
"Youtube01-Psy.csv",
"Youtube02-KatyPerry.csv",
"Youtube03-LMFAO.csv",
"Youtube04-Eminem.csv",
"Youtube05-Shakira.csv",
];
const YOUTUBE_SPAM_FILENAME: &str = "youtube_spam.csv";
const YOUTUBE_SPAM_SHA256: &str =
"f172e32ca7b4ecadb926df0c836dbe6c6485c519a47a5e7d7f719f2b3553906b";
const YOUTUBE_SPAM_DATASET_NAME: &str = "youtube_spam";
const N_SAMPLES: usize = 1_956;
const N_COLUMNS: usize = 5;
const CONTENT_COLUMN: usize = 3;
const CLASS_COLUMN: usize = 4;
#[derive(Debug)]
pub struct YoutubeSpam {
dataset: Dataset<YoutubeSpamData, DatasetError>,
}
impl YoutubeSpam {
pub fn new(storage_dir: &str) -> Self {
YoutubeSpam {
dataset: Dataset::new(storage_dir, Self::load_data),
}
}
fn load_data(dir: &str) -> Result<YoutubeSpamData, DatasetError> {
let file_path = acquire_dataset(
dir,
YOUTUBE_SPAM_FILENAME,
YOUTUBE_SPAM_DATASET_NAME,
Some(YOUTUBE_SPAM_SHA256),
|temp_path| {
download_to(
YOUTUBE_SPAM_DATA_URL,
temp_path,
Some(YOUTUBE_SPAM_ZIP_FILENAME),
)?;
unzip(&temp_path.join(YOUTUBE_SPAM_ZIP_FILENAME), temp_path)?;
let combined_path = temp_path.join(YOUTUBE_SPAM_FILENAME);
let mut combined = File::create(&combined_path)?;
for name in YOUTUBE_SPAM_SOURCE_FILENAMES {
let bytes = std::fs::read(temp_path.join(name))?;
combined.write_all(&bytes)?;
}
combined.flush()?;
Ok(combined_path)
},
)?;
let file = File::open(&file_path)?;
let mut rdr = ReaderBuilder::new().has_headers(false).from_reader(file);
let mut texts: Vec<String> = Vec::with_capacity(N_SAMPLES);
let mut labels: Vec<&'static str> = Vec::with_capacity(N_SAMPLES);
for (idx, result) in rdr.records().enumerate() {
let record =
result.map_err(|e| DatasetError::csv_read_error(YOUTUBE_SPAM_DATASET_NAME, e))?;
let line_num = idx + 1;
if record.iter().all(|f| f.is_empty()) {
continue;
}
if record.len() != N_COLUMNS {
return Err(DatasetError::invalid_column_count(
YOUTUBE_SPAM_DATASET_NAME,
N_COLUMNS,
record.len(),
line_num,
));
}
if &record[0] == "COMMENT_ID" {
continue;
}
let label = match &record[CLASS_COLUMN] {
"0" => "ham",
"1" => "spam",
other => {
return Err(DatasetError::invalid_value(
YOUTUBE_SPAM_DATASET_NAME,
"CLASS",
other,
line_num,
));
}
};
labels.push(label);
texts.push(record[CONTENT_COLUMN].to_string());
}
let n_samples = labels.len();
if n_samples == 0 {
return Err(DatasetError::empty_dataset(YOUTUBE_SPAM_DATASET_NAME));
}
let texts_array = Array1::from_vec(texts);
let labels_array = Array1::from_vec(labels);
Ok((texts_array, labels_array))
}
pub fn texts(&self) -> Result<&Array1<String>, DatasetError> {
Ok(&self.dataset.load()?.0)
}
pub fn labels(&self) -> Result<&Array1<&'static str>, DatasetError> {
Ok(&self.dataset.load()?.1)
}
pub fn data(&self) -> Result<&YoutubeSpamData, DatasetError> {
self.dataset.load()
}
pub fn get_data(&self) -> Option<&YoutubeSpamData> {
self.dataset.get()
}
pub fn get_data_mut(&mut self) -> Option<&mut YoutubeSpamData> {
self.dataset.get_mut()
}
pub fn into_data(self) -> Result<YoutubeSpamData, DatasetError> {
self.dataset.load()?;
Ok(self
.dataset
.into_inner()
.expect("data is present after a successful load"))
}
pub fn take_data(&mut self) -> Result<YoutubeSpamData, DatasetError> {
self.dataset.load()?;
Ok(self
.dataset
.take()
.expect("data is present after a successful load"))
}
}