#![warn(clippy::all, clippy::pedantic)]
use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use super::core::{TrainingDataset, TrainingSample, DatasetStats, TrainingFormat};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HfDatasetConfig {
pub source: String,
pub split: String,
pub format: Option<String>,
pub rpl_filter: Option<HashMap<String, serde_json::Value>>,
pub revision: Option<String>,
pub streaming: bool,
pub trust_remote_code: bool,
pub num_proc: Option<usize>,
}
pub struct HuggingFaceDataset {
pub name: String,
pub split: String,
pub data: Vec<serde_json::Value>,
pub features: HashMap<String, serde_json::Value>,
pub metadata: HashMap<String, serde_json::Value>,
}
impl HuggingFaceDataset {
pub async fn load(name: &str, split: &str, cache_dir: &Path) -> Result<Self> {
println!("🔄 Loading dataset {} from HuggingFace Hub...", name);
let cache = hf_hub::Cache::new(cache_dir.to_path_buf());
let repo = cache.dataset(name.to_string());
let json_files = Self::find_json_files(&repo, name).await?;
if json_files.is_empty() {
bail!("No JSON files found in dataset {}", name);
}
let data_file = &json_files[0];
println!("📁 Found data file: {}", data_file.display());
let data = Self::load_json_data(data_file).await?;
let features = Self::infer_features(&data)?;
let metadata = Self::extract_metadata(&repo).await?;
println!("✅ Successfully loaded {} samples from {}", data.len(), name);
Ok(HuggingFaceDataset {
name: name.to_string(),
split: split.to_string(),
data,
features,
metadata,
})
}
async fn find_json_files(
repo: &hf_hub::CacheRepo,
name: &str,
) -> Result<Vec<PathBuf>> {
let mut json_files = Vec::new();
let possible_files = vec![
format!("{}.json", name.replace('/', "--")), "train.json".to_string(),
"validation.json".to_string(), "test.json".to_string(), "data.json"
.to_string(),
];
for file_name in possible_files {
if let Some(path) = repo.get(&file_name) {
json_files.push(path);
}
}
if json_files.is_empty() {}
Ok(json_files)
}
async fn load_json_data(file_path: &Path) -> Result<Vec<serde_json::Value>> {
let content = tokio::fs::read_to_string(file_path)
.await
.with_context(|| {
format!("Failed to read JSON file: {}", file_path.display())
})?;
let json_value: serde_json::Value = serde_json::from_str(&content)
.with_context(|| {
format!("Failed to parse JSON from: {}", file_path.display())
})?;
match json_value {
serde_json::Value::Array(arr) => Ok(arr),
serde_json::Value::Object(obj) => {
if let Some(data) = obj.get("data") {
if let Some(arr) = data.as_array() {
return Ok(arr.clone());
}
}
if let Some(train) = obj.get("train") {
if let Some(arr) = train.as_array() {
return Ok(arr.clone());
}
}
Ok(vec![serde_json::Value::Object(obj)])
}
_ => bail!("Unsupported JSON structure in {}", file_path.display()),
}
}
fn infer_features(
data: &[serde_json::Value],
) -> Result<HashMap<String, serde_json::Value>> {
let mut features = HashMap::new();
if data.is_empty() {
return Ok(features);
}
let sample_size = std::cmp::min(10, data.len());
let samples = &data[..sample_size];
let mut key_types = HashMap::new();
for sample in samples {
if let Some(obj) = sample.as_object() {
for (key, value) in obj {
let type_str = match value {
serde_json::Value::String(_) => "string",
serde_json::Value::Number(_) => "number",
serde_json::Value::Bool(_) => "boolean",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
serde_json::Value::Null => "null",
};
key_types.insert(key.clone(), type_str.to_string());
}
}
}
for (key, type_str) in key_types {
features
.insert(
key,
serde_json::json!({ "dtype" : type_str, "_type" : "Value" }),
);
}
Ok(features)
}
async fn extract_metadata(
_repo: &hf_hub::CacheRepo,
) -> Result<HashMap<String, serde_json::Value>> {
let mut metadata = HashMap::new();
metadata.insert("dataset_name".to_string(), serde_json::json!("unknown"));
metadata.insert("split".to_string(), serde_json::json!("unknown"));
metadata.insert("num_samples".to_string(), serde_json::json!(0));
Ok(metadata)
}
pub fn get_features(&self) -> Result<Vec<String>> {
Ok(self.features.keys().map(|s| s.to_string()).collect())
}
pub fn info(&self) -> HashMap<String, serde_json::Value> {
let mut info = HashMap::new();
info.insert("name".to_string(), serde_json::json!(self.name));
info.insert("split".to_string(), serde_json::json!(self.split));
info.insert("num_samples".to_string(), serde_json::json!(self.data.len()));
info.insert("features".to_string(), serde_json::json!(self.features));
info.insert("metadata".to_string(), serde_json::json!(self.metadata));
info
}
}
#[async_trait::async_trait]
pub trait DatasetProcessor {
async fn process(&self, dataset: HuggingFaceDataset) -> Result<TrainingDataset>;
}
pub struct PreferenceProcessor;
#[async_trait::async_trait]
impl DatasetProcessor for PreferenceProcessor {
async fn process(&self, dataset: HuggingFaceDataset) -> Result<TrainingDataset> {
let mut samples = Vec::new();
for item in dataset.data {
if let Some(obj) = item.as_object() {
let sample = TrainingSample {
prompt: obj
.get("prompt")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
chosen: obj
.get("chosen")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
rejected: obj
.get("rejected")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
completion: None,
label: None,
meta: obj
.clone()
.into_iter()
.filter(|(k, _)| {
!matches!(k.as_str(), "prompt" | "chosen" | "rejected")
})
.map(|(k, v)| (k, v))
.collect(),
};
samples.push(sample);
}
}
let format = TrainingFormat::Preference {
chosen_field: "chosen".to_string(),
rejected_field: "rejected".to_string(),
};
let statistics = Self::compute_statistics(&samples);
Ok(TrainingDataset {
samples,
format,
statistics,
})
}
}
impl PreferenceProcessor {
fn compute_statistics(samples: &[TrainingSample]) -> DatasetStats {
let total_samples = samples.len();
let mut total_prompt_length = 0;
let mut prompt_count = 0;
let mut field_coverage = HashMap::new();
for sample in samples {
if sample.prompt.is_some() {
*field_coverage.entry("prompt".to_string()).or_insert(0.0) += 1.0;
total_prompt_length += sample.prompt.as_ref().unwrap().len();
prompt_count += 1;
}
if sample.chosen.is_some() {
*field_coverage.entry("chosen".to_string()).or_insert(0.0) += 1.0;
}
if sample.rejected.is_some() {
*field_coverage.entry("rejected".to_string()).or_insert(0.0) += 1.0;
}
}
for count in field_coverage.values_mut() {
*count = *count / total_samples as f64;
}
let avg_prompt_length = if prompt_count > 0 {
total_prompt_length as f64 / prompt_count as f64
} else {
0.0
};
let quality_score = Some(
(field_coverage.get("prompt").unwrap_or(&0.0)
+ field_coverage.get("chosen").unwrap_or(&0.0)
+ field_coverage.get("rejected").unwrap_or(&0.0)) / 3.0,
);
DatasetStats {
total_samples,
avg_prompt_length,
avg_completion_length: 0.0,
field_coverage,
quality_score,
}
}
}
pub struct CompletionProcessor;
#[async_trait::async_trait]
impl DatasetProcessor for CompletionProcessor {
async fn process(&self, dataset: HuggingFaceDataset) -> Result<TrainingDataset> {
let mut samples = Vec::new();
for item in dataset.data {
if let Some(obj) = item.as_object() {
let sample = TrainingSample {
prompt: obj
.get("prompt")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
chosen: None,
rejected: None,
completion: obj
.get("completion")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
label: obj.get("label").and_then(|v| v.as_f64()).map(|f| f as f32),
meta: obj
.clone()
.into_iter()
.filter(|(k, _)| {
!matches!(k.as_str(), "prompt" | "completion" | "label")
})
.map(|(k, v)| (k, v))
.collect(),
};
samples.push(sample);
}
}
let format = TrainingFormat::Completion {
completion_field: "completion".to_string(),
label_field: Some("label".to_string()),
};
let statistics = Self::compute_statistics(&samples);
Ok(TrainingDataset {
samples,
format,
statistics,
})
}
}
impl CompletionProcessor {
fn compute_statistics(samples: &[TrainingSample]) -> DatasetStats {
let total_samples = samples.len();
let mut total_prompt_length = 0;
let mut total_completion_length = 0;
let mut prompt_count = 0;
let mut completion_count = 0;
let mut field_coverage = HashMap::new();
for sample in samples {
if sample.prompt.is_some() {
*field_coverage.entry("prompt".to_string()).or_insert(0.0) += 1.0;
total_prompt_length += sample.prompt.as_ref().unwrap().len();
prompt_count += 1;
}
if sample.completion.is_some() {
*field_coverage.entry("completion".to_string()).or_insert(0.0) += 1.0;
total_completion_length += sample.completion.as_ref().unwrap().len();
completion_count += 1;
}
if sample.label.is_some() {
*field_coverage.entry("label".to_string()).or_insert(0.0) += 1.0;
}
}
for count in field_coverage.values_mut() {
*count = *count / total_samples as f64;
}
let avg_prompt_length = if prompt_count > 0 {
total_prompt_length as f64 / prompt_count as f64
} else {
0.0
};
let avg_completion_length = if completion_count > 0 {
total_completion_length as f64 / completion_count as f64
} else {
0.0
};
let quality_score = Some(
(field_coverage.get("prompt").unwrap_or(&0.0)
+ field_coverage.get("completion").unwrap_or(&0.0)) / 2.0,
);
DatasetStats {
total_samples,
avg_prompt_length,
avg_completion_length,
field_coverage,
quality_score,
}
}
}
pub struct InstructionProcessor;
#[async_trait::async_trait]
impl DatasetProcessor for InstructionProcessor {
async fn process(&self, dataset: HuggingFaceDataset) -> Result<TrainingDataset> {
let mut samples = Vec::new();
for item in dataset.data {
if let Some(obj) = item.as_object() {
let sample = TrainingSample {
prompt: obj
.get("instruction")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
chosen: None,
rejected: None,
completion: obj
.get("output")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
label: None,
meta: obj
.clone()
.into_iter()
.filter(|(k, _)| !matches!(k.as_str(), "instruction" | "output"))
.map(|(k, v)| (k, v))
.collect(),
};
samples.push(sample);
}
}
let format = TrainingFormat::Instruction {
instruction_field: "instruction".to_string(),
output_field: "output".to_string(),
};
let statistics = Self::compute_statistics(&samples);
Ok(TrainingDataset {
samples,
format,
statistics,
})
}
}
impl InstructionProcessor {
fn compute_statistics(samples: &[TrainingSample]) -> DatasetStats {
let total_samples = samples.len();
let mut total_prompt_length = 0;
let mut total_completion_length = 0;
let mut prompt_count = 0;
let mut completion_count = 0;
let mut field_coverage = HashMap::new();
for sample in samples {
if sample.prompt.is_some() {
*field_coverage.entry("instruction".to_string()).or_insert(0.0) += 1.0;
total_prompt_length += sample.prompt.as_ref().unwrap().len();
prompt_count += 1;
}
if sample.completion.is_some() {
*field_coverage.entry("output".to_string()).or_insert(0.0) += 1.0;
total_completion_length += sample.completion.as_ref().unwrap().len();
completion_count += 1;
}
}
for count in field_coverage.values_mut() {
*count = *count / total_samples as f64;
}
let avg_prompt_length = if prompt_count > 0 {
total_prompt_length as f64 / prompt_count as f64
} else {
0.0
};
let avg_completion_length = if completion_count > 0 {
total_completion_length as f64 / completion_count as f64
} else {
0.0
};
let quality_score = Some(
(field_coverage.get("instruction").unwrap_or(&0.0)
+ field_coverage.get("output").unwrap_or(&0.0)) / 2.0,
);
DatasetStats {
total_samples,
avg_prompt_length,
avg_completion_length,
field_coverage,
quality_score,
}
}
}
pub struct HfProcessor {
cache_dir: PathBuf,
processors: HashMap<String, Box<dyn DatasetProcessor + Send + Sync>>,
}
impl HfProcessor {
pub fn new(cache_dir: PathBuf) -> Self {
let mut processors: HashMap<String, Box<dyn DatasetProcessor + Send + Sync>> = HashMap::new();
processors.insert("preference".to_string(), Box::new(PreferenceProcessor));
processors.insert("completion".to_string(), Box::new(CompletionProcessor));
processors.insert("instruction".to_string(), Box::new(InstructionProcessor));
Self { cache_dir, processors }
}
pub async fn process_dataset(
&self,
dataset_name: &str,
config: &HfDatasetConfig,
) -> Result<TrainingDataset> {
let dataset = HuggingFaceDataset::load(
dataset_name,
&config.split,
&self.cache_dir,
)
.await?;
let dataset_type = self.detect_dataset_type(&dataset)?;
let processor = self
.processors
.get(&dataset_type)
.ok_or_else(|| {
anyhow::anyhow!("No processor for dataset type: {}", dataset_type)
})?;
let mut processed = processor.process(dataset).await?;
if let Some(filters) = &config.rpl_filter {
processed = self.apply_filters(processed, filters)?;
}
Ok(processed)
}
fn detect_dataset_type(&self, dataset: &HuggingFaceDataset) -> Result<String> {
let features = dataset.get_features()?;
if features.contains(&"chosen".to_string())
&& features.contains(&"rejected".to_string())
{
Ok("preference".to_string())
} else if features.contains(&"completion".to_string())
&& features.contains(&"label".to_string())
{
Ok("completion".to_string())
} else if features.contains(&"instruction".to_string())
&& features.contains(&"output".to_string())
{
Ok("instruction".to_string())
} else {
bail!("Cannot determine dataset type from features: {:?}", features)
}
}
fn apply_filters(
&self,
dataset: TrainingDataset,
_filters: &HashMap<String, serde_json::Value>,
) -> Result<TrainingDataset> {
Ok(dataset)
}
}
impl Default for HfProcessor {
fn default() -> Self {
Self::new(PathBuf::from("./hf_cache"))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_hf_processor_creation() {
let processor = HfProcessor::default();
assert!(processor.processors.contains_key("preference"));
assert!(processor.processors.contains_key("completion"));
assert!(processor.processors.contains_key("instruction"));
}
#[tokio::test]
async fn test_preference_processor() {
let mut features = HashMap::new();
features
.insert(
"prompt".to_string(),
serde_json::json!({ "dtype" : "string", "_type" : "Value" }),
);
features
.insert(
"chosen".to_string(),
serde_json::json!({ "dtype" : "string", "_type" : "Value" }),
);
features
.insert(
"rejected".to_string(),
serde_json::json!({ "dtype" : "string", "_type" : "Value" }),
);
let dataset = HuggingFaceDataset {
name: "test".to_string(),
split: "train".to_string(),
data: vec![
serde_json::json!({ "prompt" : "Test prompt", "chosen" : "Good response",
"rejected" : "Bad response" })
],
features,
metadata: HashMap::new(),
};
let processor = PreferenceProcessor;
let result = processor.process(dataset).await.unwrap();
assert_eq!(result.samples.len(), 1);
assert_eq!(result.samples[0].prompt.as_ref().unwrap(), "Test prompt");
assert_eq!(result.samples[0].chosen.as_ref().unwrap(), "Good response");
assert_eq!(result.samples[0].rejected.as_ref().unwrap(), "Bad response");
}
#[tokio::test]
async fn test_completion_processor() {
let mut features = HashMap::new();
features
.insert(
"prompt".to_string(),
serde_json::json!({ "dtype" : "string", "_type" : "Value" }),
);
features
.insert(
"completion".to_string(),
serde_json::json!({ "dtype" : "string", "_type" : "Value" }),
);
features
.insert(
"label".to_string(),
serde_json::json!({ "dtype" : "number", "_type" : "Value" }),
);
let dataset = HuggingFaceDataset {
name: "test".to_string(),
split: "train".to_string(),
data: vec![
serde_json::json!({ "prompt" : "Test prompt", "completion" :
"Test completion", "label" : 1.0 })
],
features,
metadata: HashMap::new(),
};
let processor = CompletionProcessor;
let result = processor.process(dataset).await.unwrap();
assert_eq!(result.samples.len(), 1);
assert_eq!(result.samples[0].prompt.as_ref().unwrap(), "Test prompt");
assert_eq!(result.samples[0].completion.as_ref().unwrap(), "Test completion");
assert_eq!(result.samples[0].label.unwrap(), 1.0);
}
}