mirror of https://github.com/razor-ai/soup.git
299 lines
10 KiB
Python
299 lines
10 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
|
|
"""
|
|
|
|
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"},
|
|
}
|
|
|
|
|
|
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)
|
|
# plaintext ("text" key only) checked last to avoid false matches
|
|
check_order = [
|
|
"alpaca", "llava", "sharegpt4v", "kto", "dpo", "embedding",
|
|
"audio", "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"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",
|
|
)
|
|
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)
|
|
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 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}
|