mirror of https://github.com/razor-ai/soup.git
1379 lines
44 KiB
Python
1379 lines
44 KiB
Python
"""soup data — dataset inspection and tools."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
import typer
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
|
|
from soup_cli.data.loader import load_raw_data
|
|
from soup_cli.data.validator import validate_and_stats
|
|
|
|
console = Console()
|
|
|
|
app = typer.Typer(no_args_is_help=True)
|
|
|
|
|
|
@app.command()
|
|
def inspect(
|
|
path: str = typer.Argument(..., help="Path to dataset file (jsonl, csv, parquet)"),
|
|
rows: int = typer.Option(5, "--rows", "-r", help="Number of sample rows to show"),
|
|
):
|
|
"""Inspect a dataset: show stats and sample rows."""
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
console.print(f"[dim]Inspecting {file_path}...[/]\n")
|
|
data = load_raw_data(file_path)
|
|
result = validate_and_stats(data)
|
|
|
|
# Print stats
|
|
stats_table = Table(title="Dataset Stats")
|
|
stats_table.add_column("Metric", style="bold")
|
|
stats_table.add_column("Value")
|
|
stats_table.add_row("Total samples", str(result["total"]))
|
|
stats_table.add_row("Columns", ", ".join(result["columns"]))
|
|
stats_table.add_row("Avg length (chars)", str(result["avg_length"]))
|
|
stats_table.add_row("Min length", str(result["min_length"]))
|
|
stats_table.add_row("Max length", str(result["max_length"]))
|
|
stats_table.add_row("Empty fields", str(result["empty_fields"]))
|
|
stats_table.add_row("Duplicates", str(result["duplicates"]))
|
|
console.print(stats_table)
|
|
|
|
# Vision stats (if dataset contains images)
|
|
_show_vision_stats(data)
|
|
|
|
# Print sample rows
|
|
if rows > 0 and len(data) > 0:
|
|
console.print(f"\n[bold]Sample rows ({min(rows, len(data))}):[/]")
|
|
sample_table = Table(show_lines=True)
|
|
for col in result["columns"][:5]: # max 5 columns
|
|
sample_table.add_column(col, max_width=60)
|
|
for row in data[: min(rows, len(data))]:
|
|
values = [str(row.get(col, ""))[:60] for col in result["columns"][:5]]
|
|
sample_table.add_row(*values)
|
|
console.print(sample_table)
|
|
|
|
|
|
@app.command()
|
|
def validate(
|
|
path: str = typer.Argument(..., help="Path to dataset file"),
|
|
fmt: str = typer.Option(
|
|
"auto", "--format", "-f",
|
|
help="Expected format: auto, alpaca, sharegpt, chatml, dpo, kto, plaintext",
|
|
),
|
|
):
|
|
"""Validate dataset format and report issues."""
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
data = load_raw_data(file_path)
|
|
|
|
# Auto-detect format if not specified
|
|
if fmt == "auto":
|
|
from soup_cli.data.formats import detect_format
|
|
|
|
try:
|
|
fmt = detect_format(data)
|
|
console.print(f"[dim]Auto-detected format: {fmt}[/]")
|
|
except ValueError as exc:
|
|
console.print(f"[red]{exc}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
result = validate_and_stats(data, expected_format=fmt)
|
|
|
|
if result["issues"]:
|
|
console.print("[yellow]Issues found:[/]")
|
|
for issue in result["issues"]:
|
|
console.print(f" [yellow]![/] {issue}")
|
|
else:
|
|
console.print("[bold green]Dataset is valid![/]")
|
|
|
|
valid = result["valid_rows"]
|
|
total = result["total"]
|
|
console.print(f"\n[green]{valid}/{total} rows valid for {fmt} format[/]")
|
|
|
|
|
|
@app.command()
|
|
def convert(
|
|
path: str = typer.Argument(..., help="Input dataset file"),
|
|
to: str = typer.Option(
|
|
..., "--to", "-t",
|
|
help="Target format: alpaca, sharegpt, chatml",
|
|
),
|
|
output: str = typer.Option(
|
|
None, "--output", "-o",
|
|
help="Output file path (default: <input>_<format>.jsonl)",
|
|
),
|
|
):
|
|
"""Convert a dataset between formats (alpaca, sharegpt, chatml)."""
|
|
from soup_cli.data.formats import (
|
|
CONVERTIBLE_FORMATS,
|
|
detect_format,
|
|
format_to_messages,
|
|
messages_to_format,
|
|
)
|
|
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
if to not in CONVERTIBLE_FORMATS:
|
|
console.print(
|
|
f"[red]Invalid target format: {to}[/]\n"
|
|
f"Supported: {', '.join(CONVERTIBLE_FORMATS)}"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
data = load_raw_data(file_path)
|
|
if not data:
|
|
console.print("[red]Dataset is empty.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
src_fmt = detect_format(data)
|
|
console.print(f"[dim]Detected source format: {src_fmt}[/]")
|
|
|
|
if src_fmt == to:
|
|
console.print(f"[yellow]Source and target format are both '{to}'. Nothing to convert.[/]")
|
|
raise typer.Exit()
|
|
|
|
if src_fmt == "dpo":
|
|
console.print("[red]Cannot convert DPO format (preference pairs are not conversations).[/]")
|
|
raise typer.Exit(1)
|
|
|
|
# Convert: source -> messages -> target
|
|
converted = []
|
|
failed = 0
|
|
for row in data:
|
|
messages = format_to_messages(row, src_fmt)
|
|
if messages is None:
|
|
failed += 1
|
|
continue
|
|
result = messages_to_format(messages, to)
|
|
if result is None:
|
|
failed += 1
|
|
continue
|
|
converted.append(result)
|
|
|
|
if not converted:
|
|
console.print("[red]All rows failed to convert.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
# Determine output path
|
|
if output is None:
|
|
output = str(file_path.stem) + f"_{to}.jsonl"
|
|
out_path = Path(output)
|
|
|
|
_write_jsonl(out_path, converted)
|
|
|
|
console.print(
|
|
f"[green]Converted {len(converted)} rows:[/] {src_fmt} -> {to}\n"
|
|
f"Output: [bold]{out_path}[/]"
|
|
)
|
|
if failed > 0:
|
|
console.print(f"[yellow]{failed} rows failed to convert.[/]")
|
|
|
|
|
|
@app.command()
|
|
def merge(
|
|
files: list[str] = typer.Argument(..., help="Paths to dataset files to merge"),
|
|
output: str = typer.Option(
|
|
"merged.jsonl", "--output", "-o",
|
|
help="Output file path",
|
|
),
|
|
shuffle: bool = typer.Option(False, "--shuffle", help="Shuffle after merging"),
|
|
):
|
|
"""Merge multiple datasets into a single file."""
|
|
all_data: list[dict] = []
|
|
|
|
for file_str in files:
|
|
file_path = Path(file_str)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
data = load_raw_data(file_path)
|
|
console.print(f"[dim]Loaded {len(data)} rows from {file_path}[/]")
|
|
all_data.extend(data)
|
|
|
|
if not all_data:
|
|
console.print("[red]No data loaded from any file.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
if shuffle:
|
|
random.shuffle(all_data)
|
|
|
|
out_path = Path(output)
|
|
_write_jsonl(out_path, all_data)
|
|
|
|
console.print(
|
|
f"[green]Merged {len(all_data)} rows from {len(files)} files.[/]\n"
|
|
f"Output: [bold]{out_path}[/]"
|
|
)
|
|
|
|
|
|
@app.command()
|
|
def dedup(
|
|
path: str = typer.Argument(..., help="Path to dataset file"),
|
|
output: str = typer.Option(
|
|
None, "--output", "-o",
|
|
help="Output file path (default: <input>_deduped.jsonl)",
|
|
),
|
|
threshold: float = typer.Option(
|
|
0.8, "--threshold",
|
|
help="MinHash similarity threshold (0.0-1.0)",
|
|
),
|
|
field: str = typer.Option(
|
|
None, "--field", "-f",
|
|
help="Field to hash (default: all text fields concatenated)",
|
|
),
|
|
):
|
|
"""Remove near-duplicate rows using MinHash (locality-sensitive hashing)."""
|
|
try:
|
|
from datasketch import MinHash, MinHashLSH
|
|
except ImportError:
|
|
console.print(
|
|
"[red]datasketch not installed.[/]\n"
|
|
"Install with: [bold]pip install 'soup-cli[data]'[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
data = load_raw_data(file_path)
|
|
if not data:
|
|
console.print("[red]Dataset is empty.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
console.print(f"[dim]Deduplicating {len(data)} rows (threshold={threshold})...[/]")
|
|
|
|
# Build MinHash for each row
|
|
num_perm = 128
|
|
lsh = MinHashLSH(threshold=threshold, num_perm=num_perm)
|
|
minhashes = []
|
|
|
|
for idx, row in enumerate(data):
|
|
if field:
|
|
text = str(row.get(field, ""))
|
|
else:
|
|
text = " ".join(str(v) for v in row.values() if v)
|
|
|
|
words = text.lower().split()
|
|
shingles = set()
|
|
for i in range(max(1, len(words) - 2)):
|
|
shingles.add(" ".join(words[i: i + 3]))
|
|
|
|
mhash = MinHash(num_perm=num_perm)
|
|
for shingle in shingles:
|
|
mhash.update(shingle.encode("utf-8"))
|
|
|
|
minhashes.append(mhash)
|
|
|
|
try:
|
|
lsh.insert(str(idx), mhash)
|
|
except ValueError:
|
|
pass # duplicate key, already inserted by LSH
|
|
|
|
# Collect unique indices
|
|
seen: set[int] = set()
|
|
unique_indices = []
|
|
for idx in range(len(data)):
|
|
if idx in seen:
|
|
continue
|
|
unique_indices.append(idx)
|
|
results = lsh.query(minhashes[idx])
|
|
for dup_idx_str in results:
|
|
seen.add(int(dup_idx_str))
|
|
|
|
unique_data = [data[idx] for idx in unique_indices]
|
|
removed = len(data) - len(unique_data)
|
|
|
|
# Write output
|
|
if output is None:
|
|
output = str(file_path.stem) + "_deduped.jsonl"
|
|
out_path = Path(output)
|
|
|
|
_write_jsonl(out_path, unique_data)
|
|
|
|
console.print(
|
|
f"[green]Dedup complete:[/] {len(data)} -> {len(unique_data)} rows "
|
|
f"([red]-{removed}[/] duplicates)\n"
|
|
f"Output: [bold]{out_path}[/]"
|
|
)
|
|
|
|
|
|
@app.command(name="filter")
|
|
def filter_data(
|
|
path: str = typer.Argument(..., help="Path to dataset file"),
|
|
output: str = typer.Option(
|
|
None, "--output", "-o",
|
|
help="Output file path (default: <input>_filtered.jsonl)",
|
|
),
|
|
perplexity: float = typer.Option(
|
|
None, "--perplexity", "--ppl",
|
|
help="Max perplexity threshold (rows above this are removed)",
|
|
),
|
|
coherence: float = typer.Option(
|
|
None, "--coherence",
|
|
help="Min coherence threshold 0.0-1.0 (rows below this are removed)",
|
|
),
|
|
perplexity_model: str = typer.Option(
|
|
"gpt2", "--ppl-model",
|
|
help="Model for perplexity scoring (default: gpt2)",
|
|
),
|
|
field: str = typer.Option(
|
|
None, "--field", "-f",
|
|
help="Field to score (default: all text fields concatenated)",
|
|
),
|
|
score_only: bool = typer.Option(
|
|
False, "--score-only",
|
|
help="Add scores to data without filtering (writes _scored.jsonl)",
|
|
),
|
|
):
|
|
"""Filter dataset by quality: perplexity and/or coherence scoring."""
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
if perplexity is None and coherence is None and not score_only:
|
|
console.print(
|
|
"[red]Specify at least one filter: --perplexity, --coherence, or --score-only[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
data = load_raw_data(file_path)
|
|
if not data:
|
|
console.print("[red]Dataset is empty.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
console.print(f"[dim]Scoring {len(data)} rows...[/]")
|
|
|
|
# Extract texts for scoring
|
|
texts = []
|
|
for row in data:
|
|
if field and field in row:
|
|
texts.append(str(row[field]))
|
|
else:
|
|
texts.append(" ".join(str(v) for v in row.values() if v))
|
|
|
|
# Compute coherence scores (lightweight, always computed)
|
|
from soup_cli.utils.quality import compute_coherence_score
|
|
|
|
coherence_scores = compute_coherence_score(texts)
|
|
|
|
# Compute perplexity scores (requires model, only if requested)
|
|
perplexity_scores = None
|
|
if perplexity is not None or score_only:
|
|
try:
|
|
from soup_cli.utils.quality import compute_perplexity_scores
|
|
|
|
console.print(f"[dim]Computing perplexity with {perplexity_model}...[/]")
|
|
perplexity_scores = compute_perplexity_scores(
|
|
texts, model_name=perplexity_model,
|
|
)
|
|
except ImportError:
|
|
console.print(
|
|
"[yellow]torch/transformers not available for perplexity scoring. "
|
|
"Skipping perplexity.[/]"
|
|
)
|
|
|
|
if score_only:
|
|
# Add scores to each row and write output
|
|
scored_data = []
|
|
for idx, row in enumerate(data):
|
|
scored_row = dict(row)
|
|
scored_row["_coherence_score"] = coherence_scores[idx]
|
|
if perplexity_scores is not None:
|
|
scored_row["_perplexity_score"] = round(perplexity_scores[idx], 2)
|
|
scored_data.append(scored_row)
|
|
|
|
if output is None:
|
|
output = str(file_path.stem) + "_scored.jsonl"
|
|
out_path = Path(output)
|
|
_write_jsonl(out_path, scored_data)
|
|
console.print(
|
|
f"[green]Scored {len(scored_data)} rows.[/]\n"
|
|
f"Output: [bold]{out_path}[/]"
|
|
)
|
|
return
|
|
|
|
# Filter
|
|
kept = []
|
|
removed = []
|
|
for idx, row in enumerate(data):
|
|
remove = False
|
|
if perplexity is not None and perplexity_scores is not None:
|
|
if perplexity_scores[idx] > perplexity:
|
|
remove = True
|
|
if coherence is not None and coherence_scores[idx] < coherence:
|
|
remove = True
|
|
|
|
if remove:
|
|
removed.append(row)
|
|
else:
|
|
kept.append(row)
|
|
|
|
if output is None:
|
|
output = str(file_path.stem) + "_filtered.jsonl"
|
|
out_path = Path(output)
|
|
_write_jsonl(out_path, kept)
|
|
|
|
console.print(
|
|
f"[green]Filter complete:[/] {len(data)} -> {len(kept)} rows "
|
|
f"([red]-{len(removed)}[/] removed)\n"
|
|
f"Output: [bold]{out_path}[/]"
|
|
)
|
|
if perplexity is not None and perplexity_scores is not None:
|
|
avg_ppl = sum(perplexity_scores) / len(perplexity_scores)
|
|
console.print(f"Avg perplexity: [bold]{avg_ppl:.1f}[/] (threshold: {perplexity})")
|
|
if coherence is not None:
|
|
avg_coh = sum(coherence_scores) / len(coherence_scores)
|
|
console.print(f"Avg coherence: [bold]{avg_coh:.3f}[/] (threshold: {coherence})")
|
|
|
|
|
|
@app.command()
|
|
def stats(
|
|
path: str = typer.Argument(..., help="Path to dataset file"),
|
|
):
|
|
"""Extended dataset statistics: length distribution, token counts, languages."""
|
|
from soup_cli.data.validator import extended_stats
|
|
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
data = load_raw_data(file_path)
|
|
if not data:
|
|
console.print("[red]Dataset is empty.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
ext_stats = extended_stats(data)
|
|
|
|
# Basic info table
|
|
info_table = Table(title=f"Extended Stats: {file_path.name}")
|
|
info_table.add_column("Metric", style="bold")
|
|
info_table.add_column("Value", justify="right")
|
|
info_table.add_row("Total samples", str(ext_stats["total"]))
|
|
info_table.add_row("", "")
|
|
info_table.add_row("[bold]Length (chars)[/]", "")
|
|
info_table.add_row(" p10", str(ext_stats["length_p10"]))
|
|
info_table.add_row(" p25", str(ext_stats["length_p25"]))
|
|
info_table.add_row(" p50 (median)", str(ext_stats["length_p50"]))
|
|
info_table.add_row(" p75", str(ext_stats["length_p75"]))
|
|
info_table.add_row(" p90", str(ext_stats["length_p90"]))
|
|
info_table.add_row("", "")
|
|
info_table.add_row("[bold]Tokens (approx)[/]", "")
|
|
info_table.add_row(" Average", str(ext_stats["avg_tokens"]))
|
|
info_table.add_row(" Min", str(ext_stats["min_tokens"]))
|
|
info_table.add_row(" Max", str(ext_stats["max_tokens"]))
|
|
|
|
if ext_stats["languages"]:
|
|
info_table.add_row("", "")
|
|
info_table.add_row("[bold]Languages (sample)[/]", "")
|
|
for lang, count in sorted(
|
|
ext_stats["languages"].items(), key=lambda x: -x[1]
|
|
):
|
|
info_table.add_row(f" {lang}", str(count))
|
|
|
|
console.print(info_table)
|
|
|
|
# Terminal histogram of lengths
|
|
try:
|
|
import io
|
|
import sys
|
|
|
|
import plotext as plt
|
|
|
|
lengths = ext_stats["lengths"]
|
|
if lengths:
|
|
# Force UTF-8 stdout on Windows to avoid UnicodeEncodeError
|
|
# plotext uses box-drawing chars (U+2500 etc.) that cp1251/cp1252 can't encode
|
|
original_stdout = sys.stdout
|
|
needs_redirect = (
|
|
sys.platform == "win32"
|
|
and hasattr(sys.stdout, "encoding")
|
|
and (sys.stdout.encoding or "").lower().replace("-", "") != "utf8"
|
|
)
|
|
if needs_redirect:
|
|
try:
|
|
sys.stdout = io.TextIOWrapper(
|
|
sys.stdout.buffer, encoding="utf-8", errors="replace",
|
|
)
|
|
except AttributeError:
|
|
pass # no .buffer (e.g. in tests), keep original
|
|
|
|
try:
|
|
plt.clear_figure()
|
|
plt.hist(lengths, bins=30)
|
|
plt.title("Text Length Distribution (chars)")
|
|
plt.xlabel("Length")
|
|
plt.ylabel("Count")
|
|
plt.theme("dark")
|
|
plt.show()
|
|
finally:
|
|
sys.stdout = original_stdout
|
|
except UnicodeEncodeError:
|
|
console.print(
|
|
"\n[dim]Histogram skipped (encoding issue).[/] "
|
|
"Set PYTHONIOENCODING=utf-8 to enable."
|
|
)
|
|
except ImportError:
|
|
console.print(
|
|
"\n[dim]Install plotext for histograms:[/] [bold]pip install plotext[/]"
|
|
)
|
|
|
|
|
|
def _show_vision_stats(data: list[dict]) -> None:
|
|
"""Show image statistics if dataset contains image fields."""
|
|
if not data:
|
|
return
|
|
|
|
# Check if this is a vision dataset
|
|
sample = data[0]
|
|
if "image" not in sample:
|
|
return
|
|
|
|
total = len(data)
|
|
has_image = sum(1 for row in data if row.get("image"))
|
|
missing_image = total - has_image
|
|
|
|
# Collect image file info
|
|
extensions: dict[str, int] = {}
|
|
existing = 0
|
|
for row in data:
|
|
img_path = row.get("image", "")
|
|
if not img_path:
|
|
continue
|
|
ext = Path(img_path).suffix.lower()
|
|
extensions[ext] = extensions.get(ext, 0) + 1
|
|
if Path(img_path).exists():
|
|
existing += 1
|
|
|
|
vision_table = Table(title="Vision Stats")
|
|
vision_table.add_column("Metric", style="bold")
|
|
vision_table.add_column("Value")
|
|
vision_table.add_row("Images referenced", str(has_image))
|
|
vision_table.add_row("Missing image field", str(missing_image))
|
|
vision_table.add_row("Images found on disk", str(existing))
|
|
if extensions:
|
|
ext_str = ", ".join(f"{ext} ({count})" for ext, count in sorted(extensions.items()))
|
|
vision_table.add_row("Image formats", ext_str)
|
|
console.print(vision_table)
|
|
|
|
|
|
def _write_jsonl(path: Path, data: list[dict]) -> None:
|
|
"""Write a list of dicts as JSONL."""
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
for row in data:
|
|
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
|
|
|
|
# --- Sampling strategies ---
|
|
|
|
|
|
def _sample_random(data: list[dict], num: int, seed: int | None = None) -> list[dict]:
|
|
"""Random sampling without replacement."""
|
|
rng = random.Random(seed)
|
|
num = min(num, len(data))
|
|
return rng.sample(data, num)
|
|
|
|
|
|
def _sample_diverse(
|
|
data: list[dict], num: int, seed: int | None = None
|
|
) -> list[dict]:
|
|
"""Cluster-based diverse sampling using TF-IDF + K-means.
|
|
|
|
Falls back to random sampling if sklearn is not available.
|
|
"""
|
|
num = min(num, len(data))
|
|
if num >= len(data):
|
|
return list(data)
|
|
|
|
# Extract text representations
|
|
texts = [
|
|
" ".join(str(val) for val in row.values() if val) for row in data
|
|
]
|
|
|
|
try:
|
|
from sklearn.cluster import MiniBatchKMeans
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
|
|
vectorizer = TfidfVectorizer(max_features=1000, stop_words="english")
|
|
tfidf_matrix = vectorizer.fit_transform(texts)
|
|
|
|
num_clusters = min(num, len(data))
|
|
kmeans = MiniBatchKMeans(
|
|
n_clusters=num_clusters, random_state=seed or 0, n_init=3
|
|
)
|
|
labels = kmeans.fit_predict(tfidf_matrix)
|
|
|
|
# Sample one item from each cluster (index-based dedup)
|
|
chosen_indices: list[int] = []
|
|
rng = random.Random(seed)
|
|
for cluster_id in range(num_clusters):
|
|
cluster_indices = [
|
|
idx for idx, label in enumerate(labels) if label == cluster_id
|
|
]
|
|
if cluster_indices:
|
|
chosen_indices.append(rng.choice(cluster_indices))
|
|
|
|
sampled = [data[idx] for idx in chosen_indices]
|
|
|
|
# If we need more, fill randomly from remaining
|
|
if len(sampled) < num:
|
|
remaining_indices = list(set(range(len(data))) - set(chosen_indices))
|
|
extra_indices = rng.sample(
|
|
remaining_indices, min(num - len(sampled), len(remaining_indices))
|
|
)
|
|
sampled.extend(data[idx] for idx in extra_indices)
|
|
|
|
return sampled[:num]
|
|
|
|
except ImportError:
|
|
# Fallback: simple length-based diversity (bucket by text length)
|
|
rng = random.Random(seed)
|
|
indexed = [(idx, len(texts[idx])) for idx in range(len(data))]
|
|
indexed.sort(key=lambda pair: pair[1])
|
|
# Evenly spaced picks across sorted list
|
|
step = max(1, len(indexed) // num)
|
|
picked_indices = [
|
|
indexed[idx * step][0] for idx in range(min(num, len(indexed)))
|
|
]
|
|
picked = [data[idx] for idx in picked_indices]
|
|
# Fill remainder randomly
|
|
if len(picked) < num:
|
|
remaining_indices = list(set(range(len(data))) - set(picked_indices))
|
|
extra_indices = rng.sample(
|
|
remaining_indices, min(num - len(picked), len(remaining_indices))
|
|
)
|
|
picked.extend(data[idx] for idx in extra_indices)
|
|
return picked[:num]
|
|
|
|
|
|
def _sample_hard(data: list[dict], num: int) -> list[dict]:
|
|
"""Sample hardest examples by text length (proxy for complexity).
|
|
|
|
Longer texts tend to be more complex / challenging.
|
|
"""
|
|
num = min(num, len(data))
|
|
if num >= len(data):
|
|
return list(data)
|
|
|
|
# Score by total text length (proxy for difficulty)
|
|
scored = []
|
|
for row in data:
|
|
text_len = sum(len(str(val)) for val in row.values() if val)
|
|
scored.append((text_len, row))
|
|
|
|
# Sort by length descending, take top N
|
|
scored.sort(key=lambda pair: pair[0], reverse=True)
|
|
return [row for _, row in scored[:num]]
|
|
|
|
|
|
@app.command(name="sample")
|
|
def sample_data(
|
|
path: str = typer.Argument(..., help="Path to dataset file"),
|
|
output: str = typer.Option(
|
|
None, "--output", "-o",
|
|
help="Output file path (default: <input>_sampled.jsonl)",
|
|
),
|
|
num: int = typer.Option(
|
|
None, "--n", "-n",
|
|
help="Number of samples to select",
|
|
),
|
|
pct: float = typer.Option(
|
|
None, "--pct",
|
|
help="Percentage of dataset to sample (0-100)",
|
|
),
|
|
strategy: str = typer.Option(
|
|
"random", "--strategy", "-s",
|
|
help="Sampling strategy: random, diverse (TF-IDF + clusters), hard (by length)",
|
|
),
|
|
seed: int = typer.Option(
|
|
None, "--seed",
|
|
help="Random seed for reproducibility",
|
|
),
|
|
):
|
|
"""Sample a subset of a dataset using various strategies."""
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
if num is None and pct is None:
|
|
console.print("[red]Specify either --n (count) or --pct (percentage).[/]")
|
|
raise typer.Exit(1)
|
|
|
|
if strategy not in ("random", "diverse", "hard"):
|
|
console.print(
|
|
f"[red]Unknown strategy: {strategy}[/]\n"
|
|
"Supported: [bold]random[/], [bold]diverse[/], [bold]hard[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
data = load_raw_data(file_path)
|
|
if not data:
|
|
console.print("[red]Dataset is empty.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
# Compute sample count
|
|
if pct is not None:
|
|
sample_count = max(1, int(len(data) * pct / 100))
|
|
else:
|
|
sample_count = num
|
|
|
|
# Apply strategy
|
|
if strategy == "random":
|
|
sampled = _sample_random(data, sample_count, seed=seed)
|
|
elif strategy == "diverse":
|
|
sampled = _sample_diverse(data, sample_count, seed=seed)
|
|
elif strategy == "hard":
|
|
sampled = _sample_hard(data, sample_count)
|
|
else:
|
|
sampled = _sample_random(data, sample_count, seed=seed)
|
|
|
|
# Resolve output path (with path traversal protection on explicit --output)
|
|
if output is None:
|
|
out_path = file_path.parent / f"{file_path.stem}_sampled.jsonl"
|
|
else:
|
|
out_path = Path(output).resolve()
|
|
cwd = Path.cwd().resolve()
|
|
try:
|
|
out_path.relative_to(cwd)
|
|
except ValueError:
|
|
console.print("[red]Output path must be under the current working directory.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
_write_jsonl(out_path, sampled)
|
|
|
|
console.print(
|
|
f"[green]Sampled {len(sampled)} rows[/] from {len(data)} "
|
|
f"(strategy: {strategy})\n"
|
|
f"Output: [bold]{out_path}[/]"
|
|
)
|
|
|
|
|
|
@app.command(name="split")
|
|
def split_data(
|
|
path: str = typer.Argument(..., help="Path to dataset file"),
|
|
val: int = typer.Option(
|
|
None, "--val",
|
|
help="Validation split: percentage (default) or absolute count (with --absolute)",
|
|
),
|
|
test: int = typer.Option(
|
|
None, "--test",
|
|
help="Test split: percentage (default) or absolute count (with --absolute)",
|
|
),
|
|
absolute: bool = typer.Option(
|
|
False, "--absolute",
|
|
help="Treat --val/--test as absolute sample counts instead of percentages",
|
|
),
|
|
seed: int = typer.Option(
|
|
None, "--seed",
|
|
help="Random seed for reproducible splits",
|
|
),
|
|
stratify: str = typer.Option(
|
|
None, "--stratify",
|
|
help="Field name for stratified splitting (preserves category distribution)",
|
|
),
|
|
):
|
|
"""Split dataset into train/val/test files."""
|
|
file_path = Path(path)
|
|
if not file_path.exists():
|
|
console.print(f"[red]File not found: {file_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
if val is None and test is None:
|
|
console.print("[red]Specify at least one of --val or --test.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
data = load_raw_data(file_path)
|
|
if not data:
|
|
console.print("[red]Dataset is empty.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
total = len(data)
|
|
|
|
# Calculate split sizes
|
|
if absolute:
|
|
val_count = val or 0
|
|
test_count = test or 0
|
|
if val_count + test_count >= total:
|
|
console.print(
|
|
f"[red]val ({val_count}) + test ({test_count}) >= dataset size ({total}).[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
else:
|
|
val_count = int(total * val / 100) if val else 0
|
|
test_count = int(total * test / 100) if test else 0
|
|
if val_count + test_count >= total:
|
|
console.print(
|
|
f"[red]Split sizes ({val_count} + {test_count}) >= dataset size ({total}).[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
# Perform split
|
|
if stratify:
|
|
train_data, val_data, test_data = _stratified_split(
|
|
data, val_count, test_count, stratify, seed=seed,
|
|
)
|
|
else:
|
|
train_data, val_data, test_data = _random_split(
|
|
data, val_count, test_count, seed=seed,
|
|
)
|
|
|
|
# Write output files
|
|
stem = file_path.stem
|
|
parent = file_path.parent
|
|
|
|
train_path = parent / f"{stem}_train.jsonl"
|
|
_write_jsonl(train_path, train_data)
|
|
|
|
output_msg = (
|
|
f"[green]Split {total} rows:[/]\n"
|
|
f" Train: {len(train_data)} → [bold]{train_path}[/]"
|
|
)
|
|
|
|
if val_data:
|
|
val_path = parent / f"{stem}_val.jsonl"
|
|
_write_jsonl(val_path, val_data)
|
|
output_msg += f"\n Val: {len(val_data)} → [bold]{val_path}[/]"
|
|
|
|
if test_data:
|
|
test_path = parent / f"{stem}_test.jsonl"
|
|
_write_jsonl(test_path, test_data)
|
|
output_msg += f"\n Test: {len(test_data)} → [bold]{test_path}[/]"
|
|
|
|
console.print(output_msg)
|
|
|
|
|
|
def _random_split(
|
|
data: list, val_count: int, test_count: int, seed: int | None = None,
|
|
) -> tuple:
|
|
"""Random split into train/val/test."""
|
|
rng = random.Random(seed)
|
|
indices = list(range(len(data)))
|
|
rng.shuffle(indices)
|
|
|
|
test_indices = set(indices[:test_count])
|
|
val_indices = set(indices[test_count:test_count + val_count])
|
|
|
|
train_data = []
|
|
val_data = []
|
|
test_data = []
|
|
|
|
for idx in range(len(data)):
|
|
if idx in test_indices:
|
|
test_data.append(data[idx])
|
|
elif idx in val_indices:
|
|
val_data.append(data[idx])
|
|
else:
|
|
train_data.append(data[idx])
|
|
|
|
return train_data, val_data, test_data
|
|
|
|
|
|
def _stratified_split(
|
|
data: list, val_count: int, test_count: int,
|
|
stratify_field: str, seed: int | None = None,
|
|
) -> tuple:
|
|
"""Stratified split preserving category distribution."""
|
|
# Group by stratify field
|
|
groups: dict[str, list[int]] = {}
|
|
for idx, row in enumerate(data):
|
|
key = str(row.get(stratify_field, "unknown"))
|
|
groups.setdefault(key, []).append(idx)
|
|
|
|
rng = random.Random(seed)
|
|
total = len(data)
|
|
|
|
train_indices = []
|
|
val_indices = []
|
|
test_indices = []
|
|
|
|
for key, indices in groups.items():
|
|
rng.shuffle(indices)
|
|
group_size = len(indices)
|
|
group_frac = group_size / total
|
|
|
|
group_val = round(val_count * group_frac) if val_count else 0
|
|
group_test = round(test_count * group_frac) if test_count else 0
|
|
|
|
# Ensure we don't take more than available
|
|
group_val = min(group_val, group_size)
|
|
group_test = min(group_test, group_size - group_val)
|
|
|
|
test_indices.extend(indices[:group_test])
|
|
val_indices.extend(indices[group_test:group_test + group_val])
|
|
train_indices.extend(indices[group_test + group_val:])
|
|
|
|
train_data = [data[idx] for idx in train_indices]
|
|
val_data = [data[idx] for idx in val_indices]
|
|
test_data = [data[idx] for idx in test_indices]
|
|
|
|
return train_data, val_data, test_data
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# HuggingFace Dataset Hub helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def list_datasets(search: str, sort: str = "downloads", limit: int = 20) -> list:
|
|
"""Search HuggingFace Hub for datasets. Returns list of DatasetInfo objects."""
|
|
from huggingface_hub import HfApi
|
|
|
|
api = HfApi()
|
|
return list(api.list_datasets(search=search, sort=sort, limit=limit))
|
|
|
|
|
|
def _hf_dataset_info(dataset_id: str) -> dict:
|
|
"""Fetch metadata about a HuggingFace dataset."""
|
|
from huggingface_hub import HfApi
|
|
|
|
api = HfApi()
|
|
try:
|
|
info = api.dataset_info(dataset_id)
|
|
except Exception as exc:
|
|
raise ValueError(f"Dataset not found: {dataset_id} — {exc}") from exc
|
|
|
|
# Extract split sizes
|
|
splits: dict[str, int] = {}
|
|
if hasattr(info, "card_data") and info.card_data:
|
|
ds_info = getattr(info.card_data, "dataset_info", None)
|
|
if ds_info and isinstance(ds_info, dict):
|
|
for config_data in ds_info.values():
|
|
if isinstance(config_data, dict) and "splits" in config_data:
|
|
for split_name, split_data in config_data["splits"].items():
|
|
if isinstance(split_data, dict):
|
|
splits[split_name] = split_data.get("num_examples", 0)
|
|
|
|
# Extract feature names
|
|
features: list[str] = []
|
|
if hasattr(info, "card_data") and info.card_data:
|
|
ds_info = getattr(info.card_data, "dataset_info", None)
|
|
if ds_info and isinstance(ds_info, dict):
|
|
for config_data in ds_info.values():
|
|
if isinstance(config_data, dict) and "features" in config_data:
|
|
feat_list = config_data["features"]
|
|
if isinstance(feat_list, list):
|
|
for feat in feat_list:
|
|
if isinstance(feat, dict) and "name" in feat:
|
|
features.append(feat["name"])
|
|
break
|
|
|
|
return {
|
|
"id": info.id,
|
|
"description": getattr(info, "description", "") or "",
|
|
"downloads": getattr(info, "downloads", 0) or 0,
|
|
"likes": getattr(info, "likes", 0) or 0,
|
|
"size_bytes": getattr(info, "size", None),
|
|
"splits": splits,
|
|
"features": features,
|
|
"tags": list(info.tags) if info.tags else [],
|
|
}
|
|
|
|
|
|
def _hf_download_dataset(
|
|
dataset_id: str,
|
|
split: str = "train",
|
|
samples: int | None = None,
|
|
) -> list[dict]:
|
|
"""Download a dataset from HuggingFace Hub and return as list of dicts."""
|
|
from datasets import load_dataset
|
|
|
|
try:
|
|
ds = load_dataset(
|
|
dataset_id, split=split, streaming=True, trust_remote_code=False,
|
|
)
|
|
except Exception as exc:
|
|
raise ValueError(f"Failed to load dataset {dataset_id}: {exc}") from exc
|
|
|
|
rows: list[dict] = []
|
|
for idx, row in enumerate(ds):
|
|
if samples is not None and idx >= samples:
|
|
break
|
|
rows.append(dict(row))
|
|
|
|
return rows
|
|
|
|
|
|
def _format_size_bytes(size_bytes: int | None) -> str:
|
|
"""Format byte count as human-readable string."""
|
|
if size_bytes is None:
|
|
return "unknown"
|
|
if size_bytes == 0:
|
|
return "0 B"
|
|
units = ["B", "KB", "MB", "GB", "TB"]
|
|
unit_idx = 0
|
|
size = float(size_bytes)
|
|
while size >= 1024 and unit_idx < len(units) - 1:
|
|
size /= 1024
|
|
unit_idx += 1
|
|
if unit_idx == 0:
|
|
return f"{int(size)} {units[unit_idx]}"
|
|
return f"{size:.1f} {units[unit_idx]}"
|
|
|
|
|
|
def _format_count(count: int) -> str:
|
|
"""Format large numbers with K/M suffix."""
|
|
if count >= 1_000_000:
|
|
return f"{count / 1_000_000:.1f}M"
|
|
if count >= 1_000:
|
|
return f"{count / 1_000:.1f}K"
|
|
return str(count)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# HuggingFace Dataset Hub CLI commands
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@app.command(name="search")
|
|
def search_datasets(
|
|
query: str = typer.Argument(..., help="Search query for HuggingFace datasets"),
|
|
limit: int = typer.Option(20, "--limit", "-l", help="Maximum results to show"),
|
|
sort: str = typer.Option(
|
|
"downloads", "--sort", "-s",
|
|
help="Sort by: downloads, likes, lastModified, trending, createdAt",
|
|
),
|
|
):
|
|
"""Search HuggingFace Hub for datasets."""
|
|
valid_sorts = {"downloads", "likes", "lastModified", "trending", "createdAt"}
|
|
if sort not in valid_sorts:
|
|
console.print(
|
|
f"[red]Invalid sort: {sort}[/]\n"
|
|
f"Valid options: {', '.join(sorted(valid_sorts))}"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
try:
|
|
datasets = list_datasets(search=query, sort=sort, limit=limit)
|
|
except ImportError:
|
|
console.print(
|
|
"[red]huggingface_hub not available.[/]\n"
|
|
"Install with: [bold]pip install huggingface-hub[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
except Exception as exc:
|
|
console.print(f"[red]Search failed: {exc}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
if not datasets:
|
|
console.print(f"[yellow]No datasets found for '{query}'.[/]")
|
|
return
|
|
|
|
table = Table(title=f"HuggingFace Datasets: '{query}'")
|
|
table.add_column("Dataset", style="bold cyan", max_width=45)
|
|
table.add_column("Downloads", justify="right")
|
|
table.add_column("Likes", justify="right")
|
|
table.add_column("Tags", max_width=30)
|
|
|
|
for ds_item in datasets[:limit]:
|
|
ds_tags = getattr(ds_item, "tags", []) or []
|
|
tag_str = ", ".join(ds_tags[:5])
|
|
if len(ds_tags) > 5:
|
|
tag_str += "..."
|
|
table.add_row(
|
|
ds_item.id,
|
|
_format_count(getattr(ds_item, "downloads", 0) or 0),
|
|
_format_count(getattr(ds_item, "likes", 0) or 0),
|
|
tag_str,
|
|
)
|
|
|
|
console.print(table)
|
|
console.print(f"[dim]Showing {min(limit, len(datasets))} results.[/]")
|
|
|
|
|
|
@app.command(name="preview")
|
|
def preview_dataset(
|
|
dataset_id: str = typer.Argument(
|
|
..., help="HuggingFace dataset ID (e.g. teknium/OpenHermes-2.5)"
|
|
),
|
|
):
|
|
"""Preview a remote HuggingFace dataset: metadata, splits, features."""
|
|
try:
|
|
info = _hf_dataset_info(dataset_id)
|
|
except ValueError as exc:
|
|
console.print(f"[red]{exc}[/]")
|
|
raise typer.Exit(1)
|
|
except ImportError:
|
|
console.print(
|
|
"[red]huggingface_hub not available.[/]\n"
|
|
"Install with: [bold]pip install huggingface-hub[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
table = Table(title=f"Dataset: {info['id']}")
|
|
table.add_column("Field", style="bold")
|
|
table.add_column("Value", max_width=80)
|
|
|
|
table.add_row("ID", info["id"])
|
|
desc = info["description"]
|
|
if len(desc) > 200:
|
|
desc = desc[:200] + "..."
|
|
table.add_row("Description", desc or "[dim]No description[/]")
|
|
table.add_row("Downloads", _format_count(info["downloads"]))
|
|
table.add_row("Likes", _format_count(info["likes"]))
|
|
table.add_row("Size", _format_size_bytes(info["size_bytes"]))
|
|
|
|
if info["splits"]:
|
|
splits_str = ", ".join(
|
|
f"{name} ({_format_count(count)})"
|
|
for name, count in info["splits"].items()
|
|
)
|
|
table.add_row("Splits", splits_str)
|
|
else:
|
|
table.add_row("Splits", "[dim]Not available (use streaming to explore)[/]")
|
|
|
|
if info["features"]:
|
|
table.add_row("Features", ", ".join(info["features"]))
|
|
|
|
if info["tags"]:
|
|
table.add_row("Tags", ", ".join(info["tags"][:10]))
|
|
|
|
console.print(table)
|
|
|
|
|
|
@app.command(name="download")
|
|
def download_dataset(
|
|
dataset_id: str = typer.Argument(
|
|
..., help="HuggingFace dataset ID (e.g. teknium/OpenHermes-2.5)"
|
|
),
|
|
output: str = typer.Option(
|
|
None, "--output", "-o",
|
|
help="Output file path (default: <dataset-name>.jsonl)",
|
|
),
|
|
split: str = typer.Option(
|
|
"train", "--split",
|
|
help="Dataset split to download (e.g. train, test, train[:1000])",
|
|
),
|
|
samples: int = typer.Option(
|
|
None, "--samples", "-n",
|
|
help="Max number of samples to download (streams, no full download)",
|
|
),
|
|
fmt: str = typer.Option(
|
|
None, "--format", "-f",
|
|
help="Convert to Soup format after download: alpaca, sharegpt, chatml",
|
|
),
|
|
):
|
|
"""Download a HuggingFace dataset and save as JSONL."""
|
|
max_download_samples = 1_000_000
|
|
if samples is not None and samples > max_download_samples:
|
|
console.print(
|
|
f"[red]--samples cannot exceed {max_download_samples:,}.[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
# Resolve output path
|
|
if output is None:
|
|
ds_name = dataset_id.split("/")[-1] if "/" in dataset_id else dataset_id
|
|
# Strip embedded path separators to prevent traversal
|
|
ds_name = Path(ds_name).name
|
|
out_path = (Path.cwd() / f"{ds_name}.jsonl").resolve()
|
|
cwd = Path.cwd().resolve()
|
|
try:
|
|
out_path.relative_to(cwd)
|
|
except ValueError:
|
|
console.print(
|
|
"[red]Derived output path escapes working directory.[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
else:
|
|
out_path = Path(output).resolve()
|
|
cwd = Path.cwd().resolve()
|
|
try:
|
|
out_path.relative_to(cwd)
|
|
except ValueError:
|
|
console.print(
|
|
"[red]Output path must be under the current working directory.[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
from rich.panel import Panel
|
|
|
|
console.print(Panel(
|
|
"[bold yellow]Warning:[/] Downloading this dataset may execute a "
|
|
"remote dataset loading script from HuggingFace Hub.\n\n"
|
|
"Only download datasets from sources you trust.",
|
|
title="Remote Code Warning",
|
|
border_style="yellow",
|
|
))
|
|
console.print(f"[dim]Downloading {dataset_id} (split={split})...[/]")
|
|
|
|
try:
|
|
data = _hf_download_dataset(
|
|
dataset_id, split=split, samples=samples,
|
|
)
|
|
except ValueError as exc:
|
|
console.print(f"[red]{exc}[/]")
|
|
raise typer.Exit(1)
|
|
except ImportError:
|
|
console.print(
|
|
"[red]datasets library not available.[/]\n"
|
|
"Install with: [bold]pip install datasets[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
if not data:
|
|
console.print("[red]No data downloaded (dataset may be empty).[/]")
|
|
raise typer.Exit(1)
|
|
|
|
# Optional format conversion
|
|
if fmt:
|
|
from soup_cli.data.formats import (
|
|
CONVERTIBLE_FORMATS,
|
|
detect_format,
|
|
format_to_messages,
|
|
messages_to_format,
|
|
)
|
|
|
|
if fmt not in CONVERTIBLE_FORMATS:
|
|
console.print(
|
|
f"[red]Invalid format: {fmt}[/]\n"
|
|
f"Supported: {', '.join(CONVERTIBLE_FORMATS)}"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
try:
|
|
src_fmt = detect_format(data)
|
|
except ValueError:
|
|
src_fmt = None
|
|
|
|
if src_fmt and src_fmt != fmt:
|
|
converted = []
|
|
for row in data:
|
|
messages = format_to_messages(row, src_fmt)
|
|
if messages is not None:
|
|
result = messages_to_format(messages, fmt)
|
|
if result is not None:
|
|
converted.append(result)
|
|
if converted:
|
|
data = converted
|
|
console.print(
|
|
f"[dim]Converted {len(data)} rows to {fmt} format.[/]"
|
|
)
|
|
|
|
# Apply samples limit if data came from non-streaming path
|
|
if samples is not None and len(data) > samples:
|
|
data = data[:samples]
|
|
|
|
_write_jsonl(out_path, data)
|
|
console.print(
|
|
f"[green]Downloaded {len(data)} rows.[/]\n"
|
|
f"Output: [bold]{out_path}[/]"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dataset registry CLI commands
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _get_registry_path() -> Path:
|
|
"""Get the default registry path (~/.soup/datasets.json)."""
|
|
from soup_cli.utils.registry import _default_registry_path
|
|
|
|
return _default_registry_path()
|
|
|
|
|
|
@app.command(name="register")
|
|
def register_data(
|
|
name: str = typer.Option(..., "--name", "-n", help="Dataset name"),
|
|
path: str = typer.Option(..., "--path", "-p", help="Path to dataset file"),
|
|
fmt: str = typer.Option(
|
|
"auto", "--format", "-f",
|
|
help="Dataset format: alpaca, sharegpt, chatml, dpo, kto, auto",
|
|
),
|
|
):
|
|
"""Register a local dataset by name for use in soup.yaml."""
|
|
from soup_cli.utils.registry import register_dataset
|
|
|
|
# Path traversal protection
|
|
resolved = Path(path).resolve()
|
|
cwd = Path.cwd().resolve()
|
|
try:
|
|
resolved.relative_to(cwd)
|
|
except ValueError:
|
|
console.print(
|
|
"[red]Dataset path must be under the current working directory.[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
registry_path = _get_registry_path()
|
|
|
|
try:
|
|
register_dataset(name, str(resolved), fmt, registry_path=registry_path)
|
|
except ValueError as exc:
|
|
console.print(f"[red]{exc}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
console.print(
|
|
f"[green]Registered dataset '[bold]{name}[/bold]'[/]\n"
|
|
f" Path: {path}\n"
|
|
f" Format: {fmt}"
|
|
)
|
|
|
|
|
|
@app.command(name="unregister")
|
|
def unregister_data(
|
|
name: str = typer.Option(..., "--name", "-n", help="Dataset name to remove"),
|
|
):
|
|
"""Remove a dataset from the local registry."""
|
|
from soup_cli.utils.registry import unregister_dataset
|
|
|
|
registry_path = _get_registry_path()
|
|
removed = unregister_dataset(name, registry_path=registry_path)
|
|
|
|
if removed:
|
|
console.print(f"[green]Removed dataset '{name}' from registry.[/]")
|
|
else:
|
|
console.print(f"[red]Dataset '{name}' not found in registry.[/]")
|
|
raise typer.Exit(1)
|
|
|
|
|
|
@app.command(name="registry")
|
|
def list_registry():
|
|
"""List all registered datasets."""
|
|
from soup_cli.utils.registry import load_registry
|
|
|
|
registry_path = _get_registry_path()
|
|
registry = load_registry(registry_path)
|
|
|
|
if not registry:
|
|
console.print("[yellow]No datasets registered.[/]")
|
|
console.print(
|
|
"[dim]Register with: "
|
|
"soup data register --name my-data --path data.jsonl --format alpaca[/]"
|
|
)
|
|
return
|
|
|
|
table = Table(title="Registered Datasets")
|
|
table.add_column("Name", style="bold cyan")
|
|
table.add_column("Path")
|
|
table.add_column("Format")
|
|
|
|
from rich.markup import escape
|
|
|
|
for ds_name, ds_info in sorted(registry.items()):
|
|
table.add_row(
|
|
escape(ds_name),
|
|
escape(ds_info.get("path", "")),
|
|
escape(ds_info.get("format", "")),
|
|
)
|
|
|
|
console.print(table)
|