| |
| """ |
| Convert UltraChat-style records (prompt, messages[], prompt_id) into a single |
| "text" field using the chat formatting expected by training/train_gemma_unsloth.py: |
| |
| <start_of_turn>user\n...<end_of_turn>\n<start_of_turn>model\n...<end_of_turn>\n(repeat) |
| |
| Usage: |
| python scripts/convert_ultrachat_to_text.py \ |
| --in sample_data/train_sft.jsonl \ |
| --out sample_data/train_sft_text.jsonl |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| from typing import List, TypedDict, Any, Dict, cast |
|
|
| ROLE_MAP = { |
| "user": "user", |
| "assistant": "model", |
| } |
|
|
|
|
| class Msg(TypedDict): |
| content: str |
| role: str |
|
|
|
|
| def to_chat_text(messages: List[Msg]) -> str: |
| parts: List[str] = [] |
| for m in messages: |
| role = ROLE_MAP.get(m.get("role", "user"), "user") |
| content = (m.get("content", "") or "").rstrip() |
| parts.append(f"<start_of_turn>{role}\n{content}<end_of_turn>") |
| return "\n".join(parts) + "\n" |
|
|
|
|
| def convert(in_path: str, out_path: str) -> int: |
| os.makedirs(os.path.dirname(out_path), exist_ok=True) |
| n_in = 0 |
| n_out = 0 |
| with open(in_path, "r", encoding="utf-8") as fin, open(out_path, "w", encoding="utf-8") as fout: |
| for line in fin: |
| if not line.strip(): |
| continue |
| n_in += 1 |
| obj: Any = json.loads(line) |
| raw: Any = obj.get("messages") |
| if not isinstance(raw, list) or not raw: |
| continue |
| all_dicts = True |
| for x_any in cast(List[Any], raw): |
| if not isinstance(x_any, dict): |
| all_dicts = False |
| break |
| if not all_dicts: |
| continue |
| raw_list: List[Dict[str, Any]] = cast(List[Dict[str, Any]], raw) |
| messages: List[Msg] = [] |
| for item in raw_list: |
| content_any = item.get("content") |
| role_any = item.get("role") |
| if not isinstance(content_any, str) or not isinstance(role_any, str): |
| messages = [] |
| break |
| messages.append({"content": content_any, "role": role_any}) |
| if not messages: |
| continue |
| text = to_chat_text(messages) |
| fout.write(json.dumps({"text": text}, ensure_ascii=False) + "\n") |
| n_out += 1 |
| print(f"Converted {n_out}/{n_in} records -> {out_path}") |
| return 0 |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--in", dest="in_path", required=True) |
| ap.add_argument("--out", dest="out_path", required=True) |
| args = ap.parse_args() |
| return convert(args.in_path, args.out_path) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
