mirror of https://github.com/razor-ai/soup.git
409 lines
14 KiB
Python
409 lines
14 KiB
Python
"""Dataset format detection and conversion.
|
|
|
|
Supported formats:
|
|
- alpaca: {"instruction": ..., "input": ..., "output": ...}
|
|
- sharegpt: {"conversations": [{"from": "human", "value": ...}, ...]}
|
|
- chatml: {"messages": [{"role": "user", "content": ...}, ...]}
|
|
- dpo: {"prompt": ..., "chosen": ..., "rejected": ...}
|
|
- kto: {"prompt": ..., "completion": ..., "label": true/false}
|
|
- llava: {"image": ..., "conversations": [{"from": "human", "value": ...}, ...]}
|
|
- sharegpt4v: {"image": ..., "conversations": [{"from": "human", "value": ...}, ...]}
|
|
- plaintext: {"text": "..."} — raw text for continued pre-training
|
|
- embedding: {"anchor": ..., "positive": ..., "negative": ...} — sentence embedding pairs/triplets
|
|
- audio: {"audio": ..., "messages": [...]} — audio + conversation for speech models
|
|
- tool-calling: {"messages": [...], "tools": [...], "tool_calls": [...]} — function-calling training
|
|
"""
|
|
|
|
import json
|
|
from typing import Optional
|
|
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
|
|
# Required keys per format
|
|
FORMAT_SIGNATURES = {
|
|
"alpaca": {"instruction", "output"},
|
|
"sharegpt": {"conversations"},
|
|
"chatml": {"messages"},
|
|
"dpo": {"prompt", "chosen", "rejected"},
|
|
"kto": {"prompt", "completion", "label"},
|
|
"llava": {"image", "conversations"},
|
|
"sharegpt4v": {"image", "conversations"},
|
|
"embedding": {"anchor", "positive"},
|
|
"audio": {"audio", "messages"},
|
|
"plaintext": {"text"},
|
|
"tool-calling": {"messages", "tools", "tool_calls"},
|
|
}
|
|
|
|
|
|
def detect_format(data: list[dict]) -> str:
|
|
"""Auto-detect dataset format from first few rows."""
|
|
if not data:
|
|
raise ValueError("Empty dataset - cannot detect format")
|
|
|
|
sample = data[0]
|
|
keys = set(sample.keys())
|
|
|
|
# Check more specific formats first (llava/sharegpt4v before sharegpt).
|
|
# tool-calling checked before chatml (signature is a superset of chatml).
|
|
# plaintext ("text" key only) checked last to avoid false matches.
|
|
check_order = [
|
|
"alpaca", "llava", "sharegpt4v", "kto", "dpo", "embedding",
|
|
"audio", "tool-calling", "sharegpt", "chatml", "plaintext",
|
|
]
|
|
for fmt in check_order:
|
|
required_keys = FORMAT_SIGNATURES[fmt]
|
|
if required_keys.issubset(keys):
|
|
return fmt
|
|
|
|
raise ValueError(
|
|
f"Cannot detect format. Keys found: {keys}. "
|
|
f"Expected one of: alpaca (instruction, output), "
|
|
f"sharegpt (conversations), chatml (messages), "
|
|
f"dpo (prompt, chosen, rejected), "
|
|
f"kto (prompt, completion, label), "
|
|
f"llava/sharegpt4v (image, conversations), "
|
|
f"embedding (anchor, positive), "
|
|
f"audio (audio, messages), "
|
|
f"tool-calling (messages, tools, tool_calls), "
|
|
f"plaintext (text)"
|
|
)
|
|
|
|
|
|
def format_to_messages(row: dict, fmt: str) -> Optional[dict]:
|
|
"""Convert any format to normalized structure for training.
|
|
|
|
Returns:
|
|
- SFT formats: {"messages": [{"role": ..., "content": ...}, ...]}
|
|
- Vision formats: {"messages": [...], "image": "path"}
|
|
- DPO format: {"prompt": ..., "chosen": ..., "rejected": ...}
|
|
- KTO format: {"prompt": ..., "completion": ..., "label": bool}
|
|
"""
|
|
valid_formats = (
|
|
"chatml", "alpaca", "sharegpt", "dpo", "kto", "llava", "sharegpt4v",
|
|
"plaintext", "embedding", "audio", "tool-calling",
|
|
)
|
|
if fmt not in valid_formats:
|
|
raise ValueError(f"Unknown format: {fmt}")
|
|
try:
|
|
if fmt == "chatml":
|
|
return _convert_chatml(row)
|
|
elif fmt == "alpaca":
|
|
return _convert_alpaca(row)
|
|
elif fmt == "sharegpt":
|
|
return _convert_sharegpt(row)
|
|
elif fmt == "dpo":
|
|
return _convert_dpo(row)
|
|
elif fmt == "kto":
|
|
return _convert_kto(row)
|
|
elif fmt == "plaintext":
|
|
return _convert_plaintext(row)
|
|
elif fmt == "embedding":
|
|
return _convert_embedding(row)
|
|
elif fmt == "audio":
|
|
return _convert_audio(row)
|
|
elif fmt == "tool-calling":
|
|
return _convert_tool_calling(row)
|
|
else:
|
|
return _convert_vision(row)
|
|
except (KeyError, TypeError, IndexError, ValueError):
|
|
return None
|
|
|
|
|
|
def _convert_alpaca(row: dict) -> dict:
|
|
instruction = row["instruction"]
|
|
input_text = row.get("input", "")
|
|
output = row["output"]
|
|
|
|
user_content = f"{instruction}\n{input_text}".strip() if input_text else instruction
|
|
|
|
messages = [
|
|
{"role": "user", "content": user_content},
|
|
{"role": "assistant", "content": output},
|
|
]
|
|
|
|
if row.get("system"):
|
|
messages.insert(0, {"role": "system", "content": row["system"]})
|
|
|
|
return {"messages": messages}
|
|
|
|
|
|
def _convert_sharegpt(row: dict) -> dict:
|
|
conversations = row["conversations"]
|
|
role_map = {"human": "user", "gpt": "assistant", "system": "system"}
|
|
|
|
messages = []
|
|
for turn in conversations:
|
|
role = role_map.get(turn["from"], turn["from"])
|
|
messages.append({"role": role, "content": turn["value"]})
|
|
|
|
return {"messages": messages}
|
|
|
|
|
|
def _convert_chatml(row: dict) -> dict:
|
|
# Already in the right format
|
|
return {"messages": row["messages"]}
|
|
|
|
|
|
def _convert_dpo(row: dict) -> dict:
|
|
"""Convert DPO preference row to {prompt, chosen, rejected} for trl.DPOTrainer."""
|
|
return {
|
|
"prompt": row["prompt"],
|
|
"chosen": row["chosen"],
|
|
"rejected": row["rejected"],
|
|
}
|
|
|
|
|
|
def _convert_kto(row: dict) -> dict:
|
|
"""Convert KTO row to {prompt, completion, label} for trl.KTOTrainer."""
|
|
raw_label = row["label"]
|
|
if isinstance(raw_label, str):
|
|
low = raw_label.strip().lower()
|
|
if low in ("true", "1", "yes"):
|
|
label = True
|
|
elif low in ("false", "0", "no"):
|
|
label = False
|
|
else:
|
|
raise ValueError(
|
|
f"KTO label must be true/false, got string: {raw_label!r}"
|
|
)
|
|
else:
|
|
label = bool(raw_label)
|
|
return {
|
|
"prompt": row["prompt"],
|
|
"completion": row["completion"],
|
|
"label": label,
|
|
}
|
|
|
|
|
|
def _convert_plaintext(row: dict) -> dict:
|
|
"""Convert plaintext row to {text} for continued pre-training.
|
|
|
|
Input: {"text": "raw document text..."}
|
|
Output: {"text": "raw document text..."}
|
|
"""
|
|
text = row["text"]
|
|
if not isinstance(text, str) or not text.strip():
|
|
raise ValueError("Plaintext row must have a non-empty 'text' field")
|
|
return {"text": text}
|
|
|
|
|
|
def _convert_embedding(row: dict) -> dict:
|
|
"""Convert embedding row to {anchor, positive, negative?} for embedding training.
|
|
|
|
Input: {"anchor": "query text", "positive": "similar text", "negative": "dissimilar text"}
|
|
Output: {"anchor": ..., "positive": ..., "negative": ...} (negative is optional)
|
|
"""
|
|
anchor = row["anchor"]
|
|
positive = row["positive"]
|
|
if not isinstance(anchor, str) or not anchor.strip():
|
|
raise ValueError("Embedding row must have a non-empty 'anchor' field")
|
|
if not isinstance(positive, str) or not positive.strip():
|
|
raise ValueError("Embedding row must have a non-empty 'positive' field")
|
|
result = {"anchor": anchor, "positive": positive}
|
|
negative = row.get("negative")
|
|
if isinstance(negative, str) and negative.strip():
|
|
result["negative"] = negative
|
|
return result
|
|
|
|
|
|
def _convert_audio(row: dict) -> dict:
|
|
"""Convert audio format to unified messages + audio path.
|
|
|
|
Input: {"audio": "path.wav", "messages": [{"role": "user", "content": ...}, ...]}
|
|
Output: {"messages": [...], "audio": "path.wav"}
|
|
"""
|
|
audio = row["audio"]
|
|
if not isinstance(audio, str) or not audio.strip():
|
|
raise ValueError("Audio row must have a non-empty 'audio' field")
|
|
messages = row["messages"]
|
|
if not isinstance(messages, list) or len(messages) < 1:
|
|
raise ValueError("Audio row must have a 'messages' list with at least one message")
|
|
return {"messages": messages, "audio": audio}
|
|
|
|
|
|
def _convert_vision(row: dict) -> dict:
|
|
"""Convert LLaVA / ShareGPT4V vision format to unified messages + image.
|
|
|
|
Input: {"image": "path.jpg", "conversations": [{"from": "human", "value": ...}, ...]}
|
|
Output: {"messages": [...], "image": "path.jpg"}
|
|
"""
|
|
conversations = row["conversations"]
|
|
role_map = {"human": "user", "gpt": "assistant", "system": "system"}
|
|
|
|
messages = []
|
|
for turn in conversations:
|
|
role = role_map.get(turn["from"], turn["from"])
|
|
messages.append({"role": role, "content": turn["value"]})
|
|
|
|
result = {"messages": messages, "image": row["image"]}
|
|
# Preserve optional id field
|
|
if "id" in row:
|
|
result["id"] = row["id"]
|
|
return result
|
|
|
|
|
|
def _convert_tool_calling(row: dict) -> dict:
|
|
"""Normalize tool-calling row to unified messages format.
|
|
|
|
Input:
|
|
{
|
|
"messages": [{"role": "user", "content": ...}],
|
|
"tools": [{"type": "function", "function": {...}}, ...],
|
|
"tool_calls": [{"function": {"name": ..., "arguments": "json-string"}}],
|
|
}
|
|
|
|
Output (unified format — tool schema embedded in system message,
|
|
tool_calls attached to final assistant turn):
|
|
{
|
|
"messages": [
|
|
{"role": "system", "content": "<tool schema description>"},
|
|
{"role": "user", "content": "..."},
|
|
{"role": "assistant", "content": "", "tool_calls": [...]},
|
|
]
|
|
}
|
|
|
|
Security: every tool_call's 'arguments' must be JSON-parseable. Tool schemas
|
|
must be a list of dicts. Invalid rows raise ValueError and are mapped to None
|
|
by the outer handler.
|
|
"""
|
|
tools = row["tools"]
|
|
tool_calls = row["tool_calls"]
|
|
|
|
if not isinstance(tools, list):
|
|
raise ValueError("tool-calling 'tools' must be a list")
|
|
if not isinstance(tool_calls, list):
|
|
raise ValueError("tool-calling 'tool_calls' must be a list")
|
|
|
|
for tool in tools:
|
|
if not isinstance(tool, dict):
|
|
raise ValueError("tool-calling tool entries must be dicts")
|
|
|
|
normalized_tool_calls = []
|
|
for call in tool_calls:
|
|
if not isinstance(call, dict):
|
|
raise ValueError("tool_calls entries must be dicts")
|
|
func = call.get("function")
|
|
if not isinstance(func, dict):
|
|
raise ValueError("tool_calls entry missing 'function' dict")
|
|
name = func.get("name")
|
|
if not isinstance(name, str) or not name:
|
|
raise ValueError("tool_calls 'function.name' must be a non-empty string")
|
|
args = func.get("arguments", "{}")
|
|
if isinstance(args, str):
|
|
try:
|
|
json.loads(args)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(
|
|
f"tool_calls 'arguments' must be JSON-parseable: {exc}"
|
|
) from exc
|
|
args_str = args
|
|
elif isinstance(args, dict):
|
|
args_str = json.dumps(args)
|
|
else:
|
|
raise ValueError("tool_calls 'arguments' must be str or dict")
|
|
normalized_tool_calls.append({
|
|
"function": {"name": name, "arguments": args_str},
|
|
})
|
|
|
|
original_messages = row["messages"]
|
|
if not isinstance(original_messages, list) or not original_messages:
|
|
raise ValueError("tool-calling 'messages' must be a non-empty list")
|
|
|
|
tool_schema_descriptions = []
|
|
for tool in tools:
|
|
function_def = tool.get("function", {})
|
|
tool_name = function_def.get("name", "unknown")
|
|
description = function_def.get("description", "")
|
|
params = function_def.get("parameters", {})
|
|
tool_schema_descriptions.append(
|
|
f"- {tool_name}: {description}\n parameters: {json.dumps(params)}"
|
|
)
|
|
|
|
system_content = (
|
|
"You have access to the following tools. When a tool call is needed, "
|
|
"respond with a function call in JSON.\n\n"
|
|
+ "\n".join(tool_schema_descriptions)
|
|
)
|
|
|
|
messages: list[dict] = [{"role": "system", "content": system_content}]
|
|
for msg in original_messages:
|
|
if not isinstance(msg, dict) or "role" not in msg:
|
|
raise ValueError("tool-calling messages must be dicts with 'role'")
|
|
if msg["role"] == "system":
|
|
# Merge user system message into our synthesized system content
|
|
messages[0]["content"] = msg.get("content", "") + "\n\n" + messages[0]["content"]
|
|
continue
|
|
messages.append({"role": msg["role"], "content": msg.get("content", "")})
|
|
|
|
if normalized_tool_calls:
|
|
messages.append({
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": normalized_tool_calls,
|
|
})
|
|
|
|
return {"messages": messages}
|
|
|
|
|
|
def is_vision_format(fmt: str) -> bool:
|
|
"""Check if a format is a vision/multimodal format."""
|
|
return fmt in ("llava", "sharegpt4v")
|
|
|
|
|
|
def is_audio_format(fmt: str) -> bool:
|
|
"""Check if a format is an audio/speech format."""
|
|
return fmt == "audio"
|
|
|
|
|
|
# --- Reverse conversion: messages → target format ---
|
|
|
|
CONVERTIBLE_FORMATS = ("alpaca", "sharegpt", "chatml")
|
|
|
|
|
|
def messages_to_format(row: dict, target_fmt: str) -> Optional[dict]:
|
|
"""Convert unified messages format back to a specific format.
|
|
|
|
Input: {"messages": [{"role": ..., "content": ...}, ...]}
|
|
Output: dict in target format (alpaca, sharegpt, chatml)
|
|
"""
|
|
try:
|
|
if target_fmt == "chatml":
|
|
return row # already in chatml/messages format
|
|
elif target_fmt == "alpaca":
|
|
return _to_alpaca(row["messages"])
|
|
elif target_fmt == "sharegpt":
|
|
return _to_sharegpt(row["messages"])
|
|
else:
|
|
raise ValueError(f"Cannot convert to format: {target_fmt}")
|
|
except (KeyError, TypeError, IndexError):
|
|
return None
|
|
|
|
|
|
def _to_alpaca(messages: list[dict]) -> dict:
|
|
"""Convert messages to alpaca format."""
|
|
result: dict = {"instruction": "", "input": "", "output": ""}
|
|
|
|
for msg in messages:
|
|
if msg["role"] == "system":
|
|
result["system"] = msg["content"]
|
|
elif msg["role"] == "user":
|
|
result["instruction"] = msg["content"]
|
|
elif msg["role"] == "assistant":
|
|
result["output"] = msg["content"]
|
|
|
|
return result
|
|
|
|
|
|
def _to_sharegpt(messages: list[dict]) -> dict:
|
|
"""Convert messages to sharegpt format."""
|
|
role_map = {"user": "human", "assistant": "gpt", "system": "system"}
|
|
conversations = []
|
|
for msg in messages:
|
|
conversations.append({
|
|
"from": role_map.get(msg["role"], msg["role"]),
|
|
"value": msg["content"],
|
|
})
|
|
return {"conversations": conversations}
|