use anyhow::{bail, Context, Result};
use clap::Args;
use indicatif::{ProgressBar, ProgressStyle};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::Mutex;
use walkdir::WalkDir;
#[derive(Debug, Clone, PartialEq)]
pub enum DataFormat {
Auto,
Legacy,
Molds,
Custom,
}
#[derive(Debug, Clone, PartialEq)]
pub enum TrainingFormat {
Preference { chosen_field: String, rejected_field: String },
Completion { completion_field: String, label_field: Option<String> },
Instruction { instruction_field: String, output_field: String },
Chat { messages_field: String },
Custom { fields: Vec<String> },
}
impl std::str::FromStr for DataFormat {
type Err = anyhow::Error;
fn from_str(s: &str) -> Result<Self> {
match s.to_lowercase().as_str() {
"auto" => Ok(DataFormat::Auto),
"legacy" => Ok(DataFormat::Legacy),
"molds" => Ok(DataFormat::Molds),
"custom" => Ok(DataFormat::Custom),
_ => bail!("Invalid format: {}. Must be auto, legacy, molds, or custom", s),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LegacySample {
pub x: Vec<f32>,
pub y: Vec<f32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MoldsSample {
pub module_name: String,
pub file_name: String,
pub implementation: String,
pub documentation: String,
#[serde(rename = "system_context")]
pub system_context: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingSample {
pub prompt: Option<String>,
pub chosen: Option<String>,
pub rejected: Option<String>,
pub completion: Option<String>,
pub label: Option<f32>,
pub meta: HashMap<String, Value>,
}
#[derive(Debug, Clone)]
pub struct TrainingDataset {
pub samples: Vec<TrainingSample>,
pub format: TrainingFormat,
pub statistics: DatasetStats,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatasetStats {
pub total_samples: usize,
pub avg_prompt_length: f64,
pub avg_completion_length: f64,
pub field_coverage: HashMap<String, f64>,
pub quality_score: Option<f64>,
}
#[derive(Debug)]
pub struct GenericJSONDataset {
pub data: Vec<Value>,
pub format: DataFormat,
pub schema: Option<Value>,
}
impl GenericJSONDataset {
pub fn new(
json_paths: &[PathBuf],
schema_path: Option<&Path>,
data_format: DataFormat,
) -> Result<Self> {
if json_paths.is_empty() {
bail!("No JSON files provided");
}
for path in json_paths {
if !path.is_file() {
bail!("JSON file not found: {}", path.display());
}
}
let mut raw_data = Vec::new();
for path in json_paths {
let content = fs::read_to_string(path)
.with_context(|| format!("Failed to read {}", path.display()))?;
let parsed: Value = serde_json::from_str(&content)
.with_context(|| format!("Failed to parse JSON in {}", path.display()))?;
match parsed {
Value::Array(arr) => raw_data.extend(arr),
Value::Object(obj) => raw_data.push(Value::Object(obj)),
_ => {
bail!(
"Root object in {} must be an array or object, got {}", path
.display(), parsed
)
}
}
}
if raw_data.is_empty() {
bail!("All input files are empty");
}
let format = match data_format {
DataFormat::Auto => Self::detect_format(&raw_data[0])?,
_ => data_format,
};
let schema = if let Some(schema_path) = schema_path {
let schema_content = fs::read_to_string(schema_path)
.with_context(|| {
format!("Failed to read schema {}", schema_path.display())
})?;
Some(
serde_json::from_str(&schema_content)
.with_context(|| format!("Failed to parse schema JSON"))?,
)
} else {
Self::get_builtin_schema(&format)
};
if let Some(ref schema) = schema {
Self::validate_data(&raw_data, schema, &format)?;
}
Ok(GenericJSONDataset {
data: raw_data,
format,
schema,
})
}
fn detect_format(first_sample: &Value) -> Result<DataFormat> {
if let Some(obj) = first_sample.as_object() {
if obj.contains_key("module_name") {
Ok(DataFormat::Molds)
} else if obj.contains_key("x") && obj.contains_key("y") {
Ok(DataFormat::Legacy)
} else {
Ok(DataFormat::Custom)
}
} else {
bail!("First sample is not an object - cannot auto-detect format");
}
}
fn get_builtin_schema(format: &DataFormat) -> Option<Value> {
match format {
DataFormat::Legacy => {
Some(
json!(
{ "type" : "array", "items" : { "type" : "object", "required" :
["x", "y"] } }
),
)
}
DataFormat::Molds => {
Some(
json!(
{ "type" : "array", "items" : { "type" : "object", "required" :
["module_name", "file_name", "implementation", "documentation"] }
}
),
)
}
_ => None,
}
}
fn validate_data(
data: &[Value],
_schema: &Value,
format: &DataFormat,
) -> Result<()> {
let required_keys = match format {
DataFormat::Legacy => vec!["x", "y"],
DataFormat::Molds => {
vec!["module_name", "file_name", "implementation", "documentation"]
}
_ => return Ok(()),
};
for (i, sample) in data.iter().enumerate() {
if let Some(obj) = sample.as_object() {
for key in &required_keys {
if !obj.contains_key(*key) {
bail!(
"Sample {} is missing required key '{}' for {} format", i,
key, format!("{:?}", format) .to_lowercase()
);
}
}
} else {
bail!("Sample {} is not an object", i);
}
}
Ok(())
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn get_random_sample(&self) -> Option<&Value> {
if self.data.is_empty() {
None
} else {
use rand::Rng;
let mut rng = rand::thread_rng();
let idx = rng.gen_range(0..self.data.len());
Some(&self.data[idx])
}
}
pub fn stats(&self) -> HashMap<String, Value> {
let mut stats = HashMap::new();
stats.insert("num_samples".to_string(), json!(self.len()));
stats
.insert(
"format".to_string(),
json!(format!("{:?}", self.format) .to_lowercase()),
);
stats.insert("has_schema".to_string(), json!(self.schema.is_some()));
if !self.data.is_empty() {
if let Some(obj) = self.data[0].as_object() {
stats
.insert(
"sample_keys".to_string(),
json!(obj.keys().collect::< Vec < _ >> ()),
);
}
}
stats
}
pub fn detect_training_format(&self) -> Result<TrainingFormat> {
if self.data.is_empty() {
bail!("Cannot detect training format from empty dataset");
}
let first_sample = &self.data[0];
let fields = if let Some(obj) = first_sample.as_object() {
obj.keys().map(|s| s.as_str()).collect::<Vec<_>>()
} else {
bail!("First sample is not an object");
};
if fields.contains(&"chosen") && fields.contains(&"rejected") {
return Ok(TrainingFormat::Preference {
chosen_field: "chosen".to_string(),
rejected_field: "rejected".to_string(),
});
}
if fields.contains(&"completion") {
let label_field = if fields.contains(&"label") {
Some("label".to_string())
} else {
None
};
return Ok(TrainingFormat::Completion {
completion_field: "completion".to_string(),
label_field,
});
}
if fields.contains(&"instruction") && fields.contains(&"output") {
return Ok(TrainingFormat::Instruction {
instruction_field: "instruction".to_string(),
output_field: "output".to_string(),
});
}
if fields.contains(&"messages") {
return Ok(TrainingFormat::Chat {
messages_field: "messages".to_string(),
});
}
Ok(TrainingFormat::Custom {
fields: fields.into_iter().map(|s| s.to_string()).collect(),
})
}
pub fn to_training_dataset(&self) -> Result<TrainingDataset> {
let training_format = self.detect_training_format()?;
let mut samples = Vec::new();
for (i, sample) in self.data.iter().enumerate() {
if let Some(obj) = sample.as_object() {
let training_sample = self
.convert_sample_to_training(obj, &training_format)
.with_context(|| format!("Failed to convert sample {}", i))?;
samples.push(training_sample);
} else {
bail!("Sample {} is not an object", i);
}
}
let statistics = self.compute_statistics(&samples)?;
Ok(TrainingDataset {
samples,
format: training_format,
statistics,
})
}
fn convert_sample_to_training(
&self,
obj: &serde_json::Map<String, Value>,
format: &TrainingFormat,
) -> Result<TrainingSample> {
let mut sample = TrainingSample {
prompt: None,
chosen: None,
rejected: None,
completion: None,
label: None,
meta: HashMap::new(),
};
if let Some(prompt_val) = obj.get("prompt") {
if let Some(prompt_str) = prompt_val.as_str() {
sample.prompt = Some(prompt_str.to_string());
}
}
match format {
TrainingFormat::Preference { chosen_field, rejected_field } => {
if let Some(chosen_val) = obj.get(chosen_field) {
if let Some(chosen_str) = chosen_val.as_str() {
sample.chosen = Some(chosen_str.to_string());
}
}
if let Some(rejected_val) = obj.get(rejected_field) {
if let Some(rejected_str) = rejected_val.as_str() {
sample.rejected = Some(rejected_str.to_string());
}
}
}
TrainingFormat::Completion { completion_field, label_field } => {
if let Some(completion_val) = obj.get(completion_field) {
if let Some(completion_str) = completion_val.as_str() {
sample.completion = Some(completion_str.to_string());
}
}
if let Some(label_field) = label_field {
if let Some(label_val) = obj.get(label_field) {
if let Some(label_num) = label_val.as_f64() {
sample.label = Some(label_num as f32);
} else if let Some(label_bool) = label_val.as_bool() {
sample.label = Some(if label_bool { 1.0 } else { 0.0 });
}
}
}
}
TrainingFormat::Instruction { instruction_field, output_field } => {
if let Some(instruction_val) = obj.get(instruction_field) {
if let Some(instruction_str) = instruction_val.as_str() {
sample.prompt = Some(instruction_str.to_string());
}
}
if let Some(output_val) = obj.get(output_field) {
if let Some(output_str) = output_val.as_str() {
sample.completion = Some(output_str.to_string());
}
}
}
TrainingFormat::Chat { messages_field } => {
if let Some(messages_val) = obj.get(messages_field) {
sample.meta.insert("messages".to_string(), messages_val.clone());
if let Some(messages) = messages_val.as_array() {
if let Some(first_msg) = messages.first() {
if let Some(content) = first_msg
.get("content")
.and_then(|c| c.as_str())
{
sample.prompt = Some(content.to_string());
}
}
if let Some(last_msg) = messages.last() {
if let Some(content) = last_msg
.get("content")
.and_then(|c| c.as_str())
{
sample.completion = Some(content.to_string());
}
}
}
}
}
TrainingFormat::Custom { fields } => {
for field in fields {
if let Some(value) = obj.get(field) {
sample.meta.insert(field.clone(), value.clone());
}
}
}
}
for (key, value) in obj {
if !matches!(
key.as_str(), "prompt" | "chosen" | "rejected" | "completion" | "label" |
"instruction" | "output" | "messages"
) {
sample.meta.insert(key.clone(), value.clone());
}
}
Ok(sample)
}
fn compute_statistics(&self, samples: &[TrainingSample]) -> Result<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.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;
}
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("chosen").unwrap_or(&0.0)
+ field_coverage.get("rejected").unwrap_or(&0.0)
+ field_coverage.get("completion").unwrap_or(&0.0)) / 4.0,
);
Ok(DatasetStats {
total_samples,
avg_prompt_length,
avg_completion_length,
field_coverage,
quality_score,
})
}
}
impl TrainingDataset {
pub fn to_algorithm_format(
&self,
algorithm: &str,
) -> Result<Box<dyn std::any::Any>> {
match algorithm.to_lowercase().as_str() {
"bco" => {
let bco_data = self.to_bco_format()?;
Ok(Box::new(bco_data))
}
"dpo" => {
let dpo_data = self.to_dpo_format()?;
Ok(Box::new(dpo_data))
}
"ppo" => {
let ppo_data = self.to_ppo_format()?;
Ok(Box::new(ppo_data))
}
"sft" => {
let sft_data = self.to_sft_format()?;
Ok(Box::new(sft_data))
}
_ => bail!("Unsupported algorithm: {}", algorithm),
}
}
fn to_bco_format(&self) -> Result<BCODataset> {
let mut bco_samples = Vec::new();
for sample in &self.samples {
match &self.format {
TrainingFormat::Preference { .. } => {
if let (Some(chosen), Some(rejected)) = (
&sample.chosen,
&sample.rejected,
) {
let (completion, label) = if rand::random::<bool>() {
(chosen.clone(), true)
} else {
(rejected.clone(), false)
};
bco_samples
.push(BCOSample {
prompt: sample.prompt.clone().unwrap_or_default(),
completion,
label,
});
}
}
TrainingFormat::Completion { .. } => {
if let Some(completion) = &sample.completion {
let label = sample.label.map(|l| l > 0.5).unwrap_or(true);
bco_samples
.push(BCOSample {
prompt: sample.prompt.clone().unwrap_or_default(),
completion: completion.clone(),
label,
});
}
}
_ => {
if let Some(completion) = &sample.completion {
bco_samples
.push(BCOSample {
prompt: sample.prompt.clone().unwrap_or_default(),
completion: completion.clone(),
label: sample.label.map(|l| l > 0.5).unwrap_or(true),
});
}
}
}
}
Ok(BCODataset { samples: bco_samples })
}
fn to_dpo_format(&self) -> Result<DPODataset> {
let mut dpo_samples = Vec::new();
for sample in &self.samples {
if let TrainingFormat::Preference { .. } = &self.format {
if let (Some(chosen), Some(rejected)) = (
&sample.chosen,
&sample.rejected,
) {
dpo_samples
.push(DPOSample {
prompt: sample.prompt.clone().unwrap_or_default(),
chosen: chosen.clone(),
rejected: rejected.clone(),
});
}
} else {
bail!(
"DPO format requires preference-style data (chosen/rejected fields)"
);
}
}
Ok(DPODataset { samples: dpo_samples })
}
fn to_ppo_format(&self) -> Result<PPODataset> {
let mut ppo_samples = Vec::new();
for sample in &self.samples {
if let Some(completion) = &sample.completion {
ppo_samples
.push(PPOSample {
prompt: sample.prompt.clone().unwrap_or_default(),
completion: completion.clone(),
reward: sample.label.unwrap_or(0.0),
});
}
}
Ok(PPODataset { samples: ppo_samples })
}
fn to_sft_format(&self) -> Result<SFTDataset> {
let mut sft_samples = Vec::new();
for sample in &self.samples {
if let Some(completion) = &sample.completion {
sft_samples
.push(SFTSample {
prompt: sample.prompt.clone().unwrap_or_default(),
completion: completion.clone(),
});
}
}
Ok(SFTDataset { samples: sft_samples })
}
pub fn quality_assessment(&self) -> DatasetQualityReport {
let mut report = DatasetQualityReport {
overall_score: 0.0,
issues: Vec::new(),
recommendations: Vec::new(),
};
let prompt_coverage = self
.statistics
.field_coverage
.get("prompt")
.unwrap_or(&0.0);
let completion_coverage = self
.statistics
.field_coverage
.get("completion")
.unwrap_or(&0.0);
let chosen_coverage = self
.statistics
.field_coverage
.get("chosen")
.unwrap_or(&0.0);
let rejected_coverage = self
.statistics
.field_coverage
.get("rejected")
.unwrap_or(&0.0);
match &self.format {
TrainingFormat::Preference { .. } => {
if *chosen_coverage < 0.9 {
report
.issues
.push(
format!(
"Low chosen field coverage: {:.1}%", chosen_coverage * 100.0
),
);
}
if *rejected_coverage < 0.9 {
report
.issues
.push(
format!(
"Low rejected field coverage: {:.1}%", rejected_coverage *
100.0
),
);
}
report.overall_score = (chosen_coverage + rejected_coverage
+ prompt_coverage) / 3.0;
}
TrainingFormat::Completion { .. } => {
if *completion_coverage < 0.9 {
report
.issues
.push(
format!(
"Low completion field coverage: {:.1}%", completion_coverage
* 100.0
),
);
}
report.overall_score = (completion_coverage + prompt_coverage) / 2.0;
}
_ => {
report.overall_score = self.statistics.quality_score.unwrap_or(0.0);
}
}
if self.statistics.avg_prompt_length < 10.0 {
report.issues.push("Very short average prompt length".to_string());
}
if self.statistics.avg_completion_length < 10.0 {
report.issues.push("Very short average completion length".to_string());
}
if report.issues.is_empty() {
report.recommendations.push("Dataset quality looks good!".to_string());
} else {
report
.recommendations
.push(
"Consider filtering or augmenting low-quality samples".to_string(),
);
}
report
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BCOSample {
pub prompt: String,
pub completion: String,
pub label: bool,
}
#[derive(Debug, Clone)]
pub struct BCODataset {
pub samples: Vec<BCOSample>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DPOSample {
pub prompt: String,
pub chosen: String,
pub rejected: String,
}
#[derive(Debug, Clone)]
pub struct DPODataset {
pub samples: Vec<DPOSample>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PPOSample {
pub prompt: String,
pub completion: String,
pub reward: f32,
}
#[derive(Debug, Clone)]
pub struct PPODataset {
pub samples: Vec<PPOSample>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SFTSample {
pub prompt: String,
pub completion: String,
}
#[derive(Debug, Clone)]
pub struct SFTDataset {
pub samples: Vec<SFTSample>,
}
#[derive(Debug, Clone)]
pub struct DatasetQualityReport {
pub overall_score: f64,
pub issues: Vec<String>,
pub recommendations: Vec<String>,
}
fn clean_whitespace(text: &str) -> String {
text.lines()
.map(|line| line.trim_end())
.collect::<Vec<_>>()
.join("\n")
.trim()
.to_string()
}
fn process_file(
src_path: &Path,
dst_path: &Path,
schema_path: Option<&Path>,
format_override: &DataFormat,
) -> Result<()> {
let content = fs::read_to_string(src_path)
.with_context(|| format!("Failed to read {}", src_path.display()))?;
let raw: Value = serde_json::from_str(&content)
.with_context(|| format!("Failed to parse JSON in {}", src_path.display()))?;
let _temp_dataset = GenericJSONDataset::new(
&[src_path.to_path_buf()],
schema_path,
format_override.clone(),
)?;
let cleaned = if let Value::Array(arr) = raw {
let cleaned_arr: Vec<Value> = arr
.into_iter()
.map(|mut entry| {
if let Value::Object(ref mut obj) = entry {
for (_key, value) in obj.iter_mut() {
if let Value::String(ref mut s) = value {
*s = clean_whitespace(s);
}
}
}
entry
})
.collect();
Value::Array(cleaned_arr)
} else {
raw
};
dst_path
.parent()
.map(|p| fs::create_dir_all(p))
.transpose()
.with_context(|| {
format!("Failed to create directory for {}", dst_path.display())
})?;
let cleaned_json = serde_json::to_string_pretty(&cleaned)
.with_context(|| "Failed to serialize cleaned JSON")?;
fs::write(dst_path, cleaned_json)
.with_context(|| format!("Failed to write to {}", dst_path.display()))?;
Ok(())
}
pub async fn run_multi_process_clean(
src_files: Vec<PathBuf>,
dst_root: &Path,
schema_dir: Option<&Path>,
format_override: &DataFormat,
_jobs: usize,
) -> Result<()> {
if src_files.is_empty() {
bail!("No source files provided");
}
let tasks: Vec<_> = src_files
.iter()
.map(|src| {
let dst = dst_root.join(src.file_name().unwrap());
let schema_path = schema_dir
.and_then(|dir| {
let candidate = dir
.join(format!("{}.schema.json", src.file_stem() ?.to_str() ?));
if candidate.is_file() { Some(candidate) } else { None }
});
(src.clone(), dst, schema_path)
})
.collect();
let pb = ProgressBar::new(tasks.len() as u64);
pb.set_style(
ProgressStyle::default_bar()
.template(
"{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos}/{len} ({eta}) {msg}",
)
.unwrap()
.progress_chars("#>-"),
);
pb.set_message("Cleaning & validating");
let pb = Arc::new(Mutex::new(pb));
let results: Vec<Result<()>> = tasks
.iter()
.map(|(src, dst, schema_path)| {
let result = process_file(src, dst, schema_path.as_deref(), format_override);
if let Ok(pb) = pb.try_lock() {
pb.inc(1);
}
result
})
.collect();
if let Ok(pb) = pb.try_lock() {
pb.finish_with_message("Complete");
}
let errors: Vec<_> = results
.into_iter()
.enumerate()
.filter_map(|(i, r)| r.err().map(|e| (i, e)))
.collect();
if !errors.is_empty() {
for (i, e) in &errors {
eprintln!("Error processing file {}: {}", i, e);
}
bail!("{} files failed processing", errors.len());
}
Ok(())
}
pub fn gather_json_paths(
files: &[String],
data_dirs: &[String],
) -> Result<Vec<PathBuf>> {
let mut paths = Vec::new();
for file in files {
let path = PathBuf::from(file);
if !path.is_file() {
bail!("Specified file does not exist: {}", path.display());
}
paths.push(path);
}
for dir in data_dirs {
let dir_path = PathBuf::from(dir);
if !dir_path.is_dir() {
bail!("Data directory does not exist: {}", dir_path.display());
}
for entry in WalkDir::new(&dir_path).into_iter().filter_map(|e| e.ok()) {
if entry.path().extension().map_or(false, |ext| ext == "json") {
paths.push(entry.path().to_path_buf());
}
}
}
if paths.is_empty() {
bail!("No JSON files found in specified paths");
}
paths.sort();
paths.dedup();
Ok(paths)
}
pub fn find_schema_for_file(
json_path: &Path,
schema_dir: Option<&Path>,
) -> Option<PathBuf> {
schema_dir
.and_then(|dir| {
let candidate = dir
.join(format!("{}.schema.json", json_path.file_stem() ?.to_str() ?));
if candidate.is_file() { Some(candidate) } else { None }
})
}
pub async fn run_json_cmd(args: JsonArgs) -> Result<()> {
if args.multi_process {
if args.input_folder.is_none() || args.output.is_none() {
bail!("Multi-process mode requires both --input-folder and --output");
}
let src_root = args.input_folder.unwrap();
let dst_root = args.output.unwrap();
if !src_root.is_dir() {
bail!("Input folder does not exist: {}", src_root.display());
}
let src_files: Vec<_> = WalkDir::new(&src_root)
.into_iter()
.filter_map(|e| e.ok())
.filter(|e| e.path().extension().map_or(false, |ext| ext == "json"))
.map(|e| e.path().to_path_buf())
.collect();
if src_files.is_empty() {
bail!("No JSON files found in {}", src_root.display());
}
let schema_dir = args.schema_dir.as_ref();
let format_override = args.format.clone();
println!(
"🔧 Starting multi-process clean-validate:\n source: {}\n destination: {}\n workers: {}\n format: {:?}",
src_root.display(), dst_root.display(), args.jobs, format_override
);
run_multi_process_clean(
src_files,
&dst_root,
schema_dir.map(|p| p.as_path()),
&format_override,
args.jobs,
)
.await?;
println!("✅ All files cleaned and validated successfully.");
return Ok(());
}
let json_paths = gather_json_paths(&args.file, &args.data_dir)?;
let mut all_samples = Vec::new();
for path in &json_paths {
let schema_path = find_schema_for_file(
path,
args.schema_dir.as_ref().map(|p| p.as_path()),
);
println!(
"Loading {} (schema: {})", path.file_name().unwrap().to_string_lossy(),
schema_path.as_ref().map(| p | p.file_name().unwrap().to_string_lossy())
.unwrap_or(std::borrow::Cow::Borrowed("built-in"))
);
let dataset = GenericJSONDataset::new(
&[path.clone()],
schema_path.as_ref().map(|p| p.as_path()),
args.format.clone(),
)?;
all_samples.extend(dataset.data);
}
if let Some(merge_output) = &args.merge_output {
merge_output
.parent()
.map(|p| fs::create_dir_all(p))
.transpose()
.with_context(|| {
format!("Failed to create directory for {}", merge_output.display())
})?;
let merged_json = serde_json::to_string_pretty(&all_samples)
.with_context(|| "Failed to serialize merged JSON")?;
fs::write(merge_output, merged_json)
.with_context(|| {
format!("Failed to write merged output to {}", merge_output.display())
})?;
println!(
"✅ Merged {} samples into {}", all_samples.len(), merge_output.display()
);
}
if args.show_stats {
if !json_paths.is_empty() {
let temp_dataset = GenericJSONDataset::new(
&json_paths,
args.schema_dir.as_ref().map(|p| p.as_path()),
args.format,
)?;
println!(
"\n--- Dataset statistics ------------------------------------------------"
);
for (k, v) in temp_dataset.stats() {
println!("{:20}: {}", k, v);
}
println!(
"--------------------------------------------------------------------"
);
}
}
if args.merge_output.is_none() {
println!("🎉 Validation finished – no merged output requested.");
}
Ok(())
}
#[derive(Args, Debug, Clone)]
pub struct JsonArgs {
#[arg(long)]
pub data_dir: Vec<String>,
#[arg(long, short = 'f')]
pub file: Vec<String>,
#[arg(long)]
pub schema_dir: Option<PathBuf>,
#[arg(long, default_value = "auto")]
pub format: DataFormat,
#[arg(long)]
pub merge_output: Option<PathBuf>,
#[arg(long)]
pub show_stats: bool,
#[arg(long, default_value = "42")]
pub seed: u64,
#[arg(long)]
pub multi_process: bool,
#[arg(long)]
pub input_folder: Option<PathBuf>,
#[arg(long, short = 'o')]
pub output: Option<PathBuf>,
#[arg(long, short = 'j', default_value_t = num_cpus::get())]
pub jobs: usize,
}