87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_INPUT = ROOT / "plugins" / "1c" / "training" / "examples" / "instruction.examples.jsonl"
|
|
DEFAULT_GENERATED_INPUT = ROOT / "plugins" / "1c" / "training" / "raw" / "generated.instruction.jsonl"
|
|
DEFAULT_OUTPUT = ROOT / "plugins" / "1c" / "training" / "prepared" / "train.chat.jsonl"
|
|
|
|
|
|
def iter_jsonl(path: Path) -> list[dict]:
|
|
records: list[dict] = []
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
line = line.strip()
|
|
if line:
|
|
records.append(json.loads(line))
|
|
return records
|
|
|
|
|
|
def normalize_record(record: dict) -> dict:
|
|
messages = []
|
|
for message in record.get("messages") or []:
|
|
messages.append(
|
|
{
|
|
"role": message["role"],
|
|
"content": message["content"].strip(),
|
|
}
|
|
)
|
|
return {
|
|
"messages": messages,
|
|
"metadata": {
|
|
"id": record.get("id"),
|
|
**(record.get("metadata") or {}),
|
|
},
|
|
}
|
|
|
|
|
|
def dedupe_records(records: list[dict]) -> list[dict]:
|
|
seen: set[str] = set()
|
|
result = []
|
|
for record in records:
|
|
record_id = str(record.get("id") or (record.get("metadata") or {}).get("id") or "")
|
|
if record_id and record_id in seen:
|
|
continue
|
|
if record_id:
|
|
seen.add(record_id)
|
|
result.append(record)
|
|
return result
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Prepare 1C training data in chat JSONL format.")
|
|
parser.add_argument("--input", type=Path, action="append", help="Input JSONL file. Can be passed multiple times.")
|
|
parser.add_argument(
|
|
"--include-generated",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Include the generated local training JSONL when it exists.",
|
|
)
|
|
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
|
args = parser.parse_args()
|
|
|
|
input_paths = args.input or [DEFAULT_INPUT]
|
|
if args.include_generated and DEFAULT_GENERATED_INPUT.exists() and DEFAULT_GENERATED_INPUT not in input_paths:
|
|
input_paths.append(DEFAULT_GENERATED_INPUT)
|
|
|
|
raw_records = []
|
|
for input_path in input_paths:
|
|
raw_records.extend(iter_jsonl(input_path))
|
|
|
|
records = [normalize_record(record) for record in dedupe_records(raw_records)]
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with args.output.open("w", encoding="utf-8") as handle:
|
|
for record in records:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
print(f"Wrote {len(records)} training record(s) to {args.output}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|