File size: 1,418 Bytes
c676a94 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
#!/usr/bin/env python3
"""
Download, transform LMSYS-Chat-1M into plain text for LLM completion models
in the format:
<|im_start|>role
message<|endoftext|>
<|im_stop|>
with 6 newlines between conversations.
"""
from datasets import load_dataset
import sys
def main(output_path="lmsys_chat_1m.txt", split="train"):
ds = load_dataset("lmsys/lmsys-chat-1m", split=split)
with open(output_path, "w", encoding="utf-8") as out:
for i, sample in enumerate(ds):
conv = sample["conversation"] # list of messages
for msg in conv:
role = msg["role"]
content = msg["content"].strip()
out.write(f"<|im_start|>{role}\n{content}<|endoftext|>\n<|im_stop|>\n")
out.write("\n" * 6) # 6 newlines between conversations
if (i + 1) % 10000 == 0:
print(f"Processed {i + 1} conversations", file=sys.stderr)
print(f"✔ Saved plain-text to: {output_path}")
if __name__ == "__main__":
import argparse
p = argparse.ArgumentParser(description="Convert LMSYS-Chat-1M to LLM-friendly text format")
p.add_argument("--output", "-o", default="lmsys_chat_1m.txt", help="Output file path")
p.add_argument("--split", "-s", default="train", help="Dataset split (e.g. 'train')")
args = p.parse_args()
main(output_path=args.output, split=args.split)
|