use std::path::PathBuf;
#[derive(Debug)]
pub enum DatasetAction {
Process {
files: Vec<PathBuf>,
output: Option<PathBuf>,
format: Option<String>,
algorithm: Option<String>,
validate: bool,
},
Analyze {
files: Vec<PathBuf>,
detailed: bool
},
Convert {
input: PathBuf,
output: Option<PathBuf>,
from_format: String,
to_format: String,
},
Quality {
files: Vec<PathBuf>,
report: bool
},
Huggingface {
dataset: String,
split: Option<String>,
output: Option<PathBuf>,
cache_dir: Option<PathBuf>,
},
}
pub async fn dataset_command(
action: DatasetAction,
verbose: bool,
) -> Result<(), Box<dyn std::error::Error>> {
match action {
DatasetAction::Process { files, output, format, algorithm, validate } => {
if verbose {
println!("🧠 Processing datasets with HLX-AI...");
println!(" Files: {:?}", files);
println!(" Output: {:?}", output);
println!(" Format: {:?}", format);
println!(" Algorithm: {:?}", algorithm);
println!(" Validate: {}", validate);
}
use crate::map::core::{GenericJSONDataset, DataFormat};
for file in &files {
if verbose {
println!("📊 Processing: {}", file.display());
}
let dataset = GenericJSONDataset::new(
&[file.clone()],
None,
DataFormat::Auto,
)
.map_err(|e| {
format!("Failed to load dataset {}: {}", file.display(), e)
})?;
let training_dataset = dataset
.to_training_dataset()
.map_err(|e| {
format!("Failed to convert dataset {}: {}", file.display(), e)
})?;
if validate {
let quality = training_dataset.quality_assessment();
println!("✅ Quality Score: {:.2}", quality.overall_score);
if !quality.issues.is_empty() {
println!("⚠️ Issues:");
for issue in &quality.issues {
println!(" - {}", issue);
}
}
}
if let Some(algo) = &algorithm {
if training_dataset.to_algorithm_format(algo).is_ok() {
println!("✅ Converted to {} format", algo.to_uppercase());
} else {
println!(
"❌ Failed to convert to {} format", algo.to_uppercase()
);
}
}
println!(
"📈 Dataset stats: {} samples", training_dataset.samples.len()
);
}
println!("🎉 Dataset processing completed!");
Ok(())
}
DatasetAction::Analyze { files, detailed } => {
if verbose {
println!("🔍 Analyzing datasets...");
}
use crate::map::core::{GenericJSONDataset, DataFormat};
for file in files {
if verbose {
println!("📊 Analyzing: {}", file.display());
}
let dataset = GenericJSONDataset::new(
&[file.clone()],
None,
DataFormat::Auto,
)
.map_err(|e| {
format!("Failed to load dataset {}: {}", file.display(), e)
})?;
println!("\n--- Dataset Analysis: {} ---", file.display());
for (key, value) in dataset.stats() {
println!("{:15}: {}", key, value);
}
if detailed {
let training_dataset = dataset
.to_training_dataset()
.map_err(|e| {
format!(
"Failed to convert dataset {}: {}", file.display(), e
)
})?;
println!("\n--- Training Format Analysis ---");
println!("Format: {:?}", training_dataset.format);
println!("Samples: {}", training_dataset.samples.len());
println!(
"Avg Prompt Length: {:.1}", training_dataset.statistics
.avg_prompt_length
);
println!(
"Avg Completion Length: {:.1}", training_dataset.statistics
.avg_completion_length
);
println!("\n--- Field Coverage ---");
for (field, coverage) in &training_dataset.statistics.field_coverage
{
println!("{:12}: {:.1}%", field, coverage * 100.0);
}
}
}
Ok(())
}
DatasetAction::Convert { input, output: _output, from_format, to_format } => {
if verbose {
println!("🔄 Converting dataset format...");
println!(" Input: {}", input.display());
println!(" From: {}", from_format);
println!(" To: {}", to_format);
}
println!("🔄 Format conversion: {} → {}", from_format, to_format);
println!("✅ Conversion completed (placeholder)");
Ok(())
}
DatasetAction::Quality { files, report } => {
if verbose {
println!("📊 Assessing dataset quality...");
}
use crate::map::core::{GenericJSONDataset, DataFormat};
for file in files {
let dataset = GenericJSONDataset::new(
&[file.clone()],
None,
DataFormat::Auto,
)
.map_err(|e| {
format!("Failed to load dataset {}: {}", file.display(), e)
})?;
let training_dataset = dataset
.to_training_dataset()
.map_err(|e| {
format!("Failed to convert dataset {}: {}", file.display(), e)
})?;
let quality = training_dataset.quality_assessment();
if report {
println!("\n=== Quality Report: {} ===", file.display());
println!("Overall Score: {:.2}/1.0", quality.overall_score);
println!("\nIssues:");
if quality.issues.is_empty() {
println!(" ✅ No issues found");
} else {
for issue in &quality.issues {
println!(" ⚠️ {}", issue);
}
}
println!("\nRecommendations:");
for rec in &quality.recommendations {
println!(" 💡 {}", rec);
}
} else {
println!(
"📊 {}: Quality Score {:.2}", file.display(), quality
.overall_score
);
}
}
Ok(())
}
DatasetAction::Huggingface { dataset, split, output, cache_dir } => {
if verbose {
println!("🤗 Loading HuggingFace dataset...");
println!(" Dataset: {}", dataset);
println!(
" Split: {:?}", split.as_ref().unwrap_or(& "train".to_string())
);
println!(" Cache: {:?}", cache_dir);
println!(" Output: {:?}", output);
}
let processor = crate::map::HfProcessor::new(
cache_dir.unwrap_or_else(|| PathBuf::from("./hf_cache")),
);
let config = crate::map::HfDatasetConfig {
source: dataset.clone(),
split: split.unwrap_or_else(|| "train".to_string()),
format: None,
rpl_filter: None,
revision: None,
streaming: false,
trust_remote_code: false,
num_proc: None,
};
match processor.process_dataset(&dataset, &config).await {
Ok(training_dataset) => {
println!("✅ HuggingFace dataset loaded successfully");
println!("📊 Samples: {}", training_dataset.samples.len());
println!("📝 Format: {:?}", training_dataset.format);
if let Some(output_path) = output {
let json_output = serde_json::to_string_pretty(
&training_dataset.samples,
)
.map_err(|e| format!("Failed to serialize output: {}", e))?;
std::fs::write(&output_path, json_output)
.map_err(|e| {
format!(
"Failed to write output file {}: {}", output_path.display(),
e
)
})?;
println!(
"💾 Saved processed dataset to: {}", output_path.display()
);
}
}
Err(e) => {
println!("❌ Failed to load HuggingFace dataset: {}", e);
return Err(e.into());
}
}
Ok(())
}
}
}