mirror of https://github.com/razor-ai/soup.git
230 lines
7.3 KiB
Python
230 lines
7.3 KiB
Python
"""Tests for data tools: convert, merge, dedup, stats, and reverse format conversion."""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
from soup_cli.data.formats import (
|
|
format_to_messages,
|
|
messages_to_format,
|
|
)
|
|
from soup_cli.data.validator import extended_stats
|
|
|
|
runner = CliRunner()
|
|
|
|
|
|
# --- Reverse format conversion tests ---
|
|
|
|
|
|
def test_messages_to_alpaca():
|
|
"""Convert messages → alpaca format."""
|
|
row = {"messages": [
|
|
{"role": "user", "content": "What is AI?"},
|
|
{"role": "assistant", "content": "AI is artificial intelligence."},
|
|
]}
|
|
result = messages_to_format(row, "alpaca")
|
|
assert result["instruction"] == "What is AI?"
|
|
assert result["output"] == "AI is artificial intelligence."
|
|
|
|
|
|
def test_messages_to_alpaca_with_system():
|
|
"""System message should be preserved in alpaca format."""
|
|
row = {"messages": [
|
|
{"role": "system", "content": "You are helpful."},
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]}
|
|
result = messages_to_format(row, "alpaca")
|
|
assert result["system"] == "You are helpful."
|
|
assert result["instruction"] == "Hello"
|
|
assert result["output"] == "Hi!"
|
|
|
|
|
|
def test_messages_to_sharegpt():
|
|
"""Convert messages → sharegpt format."""
|
|
row = {"messages": [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]}
|
|
result = messages_to_format(row, "sharegpt")
|
|
assert result["conversations"][0]["from"] == "human"
|
|
assert result["conversations"][0]["value"] == "Hello"
|
|
assert result["conversations"][1]["from"] == "gpt"
|
|
assert result["conversations"][1]["value"] == "Hi!"
|
|
|
|
|
|
def test_messages_to_chatml():
|
|
"""Convert messages → chatml should be identity."""
|
|
row = {"messages": [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi!"},
|
|
]}
|
|
result = messages_to_format(row, "chatml")
|
|
assert result == row
|
|
|
|
|
|
def test_roundtrip_alpaca_to_chatml_to_alpaca():
|
|
"""Converting alpaca → chatml → alpaca should preserve content."""
|
|
original = {"instruction": "Explain X", "input": "", "output": "X is amazing."}
|
|
chatml = format_to_messages(original, "alpaca")
|
|
assert chatml is not None
|
|
back = messages_to_format(chatml, "alpaca")
|
|
assert back["instruction"] == original["instruction"]
|
|
assert back["output"] == original["output"]
|
|
|
|
|
|
def test_roundtrip_sharegpt_to_chatml_to_sharegpt():
|
|
"""Converting sharegpt → chatml → sharegpt should preserve content."""
|
|
original = {"conversations": [
|
|
{"from": "human", "value": "Hello"},
|
|
{"from": "gpt", "value": "Hi!"},
|
|
]}
|
|
chatml = format_to_messages(original, "sharegpt")
|
|
assert chatml is not None
|
|
back = messages_to_format(chatml, "sharegpt")
|
|
assert back["conversations"][0]["from"] == "human"
|
|
assert back["conversations"][0]["value"] == "Hello"
|
|
|
|
|
|
def test_messages_to_invalid_format():
|
|
"""Invalid target format should raise ValueError."""
|
|
row = {"messages": [{"role": "user", "content": "Hello"}]}
|
|
with pytest.raises(ValueError, match="Cannot convert"):
|
|
messages_to_format(row, "dpo")
|
|
|
|
|
|
# --- Extended stats tests ---
|
|
|
|
|
|
def test_extended_stats_basic():
|
|
"""Extended stats should compute percentiles and token counts."""
|
|
data = [
|
|
{"instruction": "short", "output": "ok"},
|
|
{"instruction": "a bit longer instruction here", "output": "response text here"},
|
|
{"instruction": "x" * 200, "output": "y" * 100},
|
|
]
|
|
result = extended_stats(data)
|
|
assert result["total"] == 3
|
|
assert result["length_p50"] > 0
|
|
assert result["avg_tokens"] > 0
|
|
assert result["min_tokens"] >= 1
|
|
assert len(result["lengths"]) == 3
|
|
|
|
|
|
def test_extended_stats_empty():
|
|
"""Extended stats on empty data should return zeros."""
|
|
result = extended_stats([])
|
|
assert result["total"] == 0
|
|
assert result["avg_tokens"] == 0
|
|
|
|
|
|
# --- CLI data command tests ---
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_alpaca_file(tmp_path: Path) -> Path:
|
|
"""Create a sample alpaca JSONL file."""
|
|
data = [
|
|
{"instruction": "What is AI?", "input": "", "output": "AI is..."},
|
|
{"instruction": "Explain ML", "input": "", "output": "ML is..."},
|
|
{"instruction": "What is DL?", "input": "", "output": "DL is..."},
|
|
]
|
|
filepath = tmp_path / "sample.jsonl"
|
|
with open(filepath, "w") as f:
|
|
for row in data:
|
|
f.write(json.dumps(row) + "\n")
|
|
return filepath
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_sharegpt_file(tmp_path: Path) -> Path:
|
|
"""Create a sample sharegpt JSONL file."""
|
|
data = [
|
|
{"conversations": [
|
|
{"from": "human", "value": "Hello"},
|
|
{"from": "gpt", "value": "Hi!"},
|
|
]},
|
|
{"conversations": [
|
|
{"from": "human", "value": "How are you?"},
|
|
{"from": "gpt", "value": "I'm good!"},
|
|
]},
|
|
]
|
|
filepath = tmp_path / "sharegpt.jsonl"
|
|
with open(filepath, "w") as f:
|
|
for row in data:
|
|
f.write(json.dumps(row) + "\n")
|
|
return filepath
|
|
|
|
|
|
def test_convert_command(sample_alpaca_file, tmp_path):
|
|
"""soup data convert should convert alpaca → sharegpt."""
|
|
output = str(tmp_path / "output.jsonl")
|
|
result = runner.invoke(
|
|
app, ["data", "convert", str(sample_alpaca_file), "--to", "sharegpt", "-o", output]
|
|
)
|
|
assert result.exit_code == 0
|
|
assert "Converted" in result.output
|
|
|
|
# Verify output file
|
|
with open(output) as f:
|
|
lines = f.readlines()
|
|
assert len(lines) == 3
|
|
first = json.loads(lines[0])
|
|
assert "conversations" in first
|
|
|
|
|
|
def test_convert_same_format(sample_alpaca_file):
|
|
"""Converting to same format should warn and exit."""
|
|
result = runner.invoke(
|
|
app, ["data", "convert", str(sample_alpaca_file), "--to", "alpaca"]
|
|
)
|
|
assert result.exit_code == 0
|
|
assert "Nothing to convert" in result.output
|
|
|
|
|
|
def test_convert_missing_file():
|
|
"""Converting nonexistent file should fail."""
|
|
result = runner.invoke(app, ["data", "convert", "nonexistent.jsonl", "--to", "chatml"])
|
|
assert result.exit_code == 1
|
|
|
|
|
|
def test_merge_command(sample_alpaca_file, sample_sharegpt_file, tmp_path):
|
|
"""soup data merge should merge files."""
|
|
output = str(tmp_path / "merged.jsonl")
|
|
result = runner.invoke(
|
|
app,
|
|
["data", "merge", str(sample_alpaca_file), str(sample_sharegpt_file), "-o", output],
|
|
)
|
|
assert result.exit_code == 0
|
|
assert "Merged" in result.output
|
|
assert "5 rows" in result.output # 3 + 2
|
|
|
|
with open(output) as f:
|
|
lines = f.readlines()
|
|
assert len(lines) == 5
|
|
|
|
|
|
def test_merge_missing_file(sample_alpaca_file):
|
|
"""Merge with nonexistent file should fail."""
|
|
result = runner.invoke(
|
|
app, ["data", "merge", str(sample_alpaca_file), "nonexistent.jsonl"]
|
|
)
|
|
assert result.exit_code == 1
|
|
|
|
|
|
def test_stats_command(sample_alpaca_file):
|
|
"""soup data stats should show extended statistics."""
|
|
result = runner.invoke(app, ["data", "stats", str(sample_alpaca_file)])
|
|
assert result.exit_code == 0
|
|
assert "p50" in result.output
|
|
assert "Tokens" in result.output
|
|
|
|
|
|
def test_stats_missing_file():
|
|
"""Stats on nonexistent file should fail."""
|
|
result = runner.invoke(app, ["data", "stats", "nonexistent.jsonl"])
|
|
assert result.exit_code == 1
|