tunesmol.py

from datasets import load_dataset
from collections import defaultdict

# ======================
# Load dataset
# ======================
dataset = load_dataset(
    "simonko912/chan-shitpost-3",
    split="train",
    streaming=True
)

# ======================
# Group by thread
# ======================
threads = defaultdict(list)

for ex in dataset:
    if ex["text"] and len(ex["text"].strip()) > 0:
        threads[ex["thread_id"]].append(ex)

# ======================
# Build SmolLM2 ChatML
# ======================
def build_chatml(posts, max_messages=50):
    posts = sorted(posts, key=lambda x: x["position_in_thread"])

    chat = ""

    for i, post in enumerate(posts[:max_messages]):
        role = "user" if i % 2 == 0 else "assistant"
        content = post["text"].strip()

        if len(content) < 2:
            continue

        chat += f"<|im_start|>{role}\n{content}<|im_end|>\n"

    # IMPORTANT: leave open assistant for training completion
    chat += "<|im_start|>assistant\n"

    return {"text": chat}

# ======================
# Convert all threads
# ======================
train_data = []

for thread_id, posts in threads.items():
    if len(posts) < 2:
        continue

    train_data.append(build_chatml(posts))

print("Conversations built:", len(train_data))