re-implement dataset filtering in more efficient way

train_with_unsloth.py

@@ -26,6 +26,8 @@ from transformers.integrations import WandbCallback
 from datasets import load_dataset
 import os
 import torch
+from datasets import load_dataset, Dataset
+import random
 
 # Project and dataset settings
 PROJECT_NAME = 'SmolLM2-360M-Instruct-TaiwanChat'
@@ -66,23 +68,40 @@ model = FastLanguageModel.get_peft_model(
 )
 
 # Prepare dataset with 5% validation split
-dataset = load_dataset(DATASET_ID, split=f"train")
-def fmt(examples):
-    return {"text": [
-        tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
-        for msgs in examples["messages"]
-    ]}
-dataset = dataset.map(fmt, batched=True, remove_columns=["messages"])
-def is_within_max_len(example):
-    toks = tokenizer(
-        example["text"], 
-        add_special_tokens=False
-    )["input_ids"]
-    return len(toks) <= MAX_LEN
-
-# Filter out samples whose encoded length >= MAX_LEN 
-filtered_ds = dataset.filter(is_within_max_len)
-dataset = filtered_ds.select(range(N_SAMPLES))
+def load_fitting_samples(dataset_id, tokenizer, max_len, n_samples, seed=3407):
+    # 1) Open the HF dataset in streaming mode
+    stream = load_dataset(dataset_id, split="train", streaming=True)
+
+    selected = []
+    for example in stream:
+        # 2) Render the chat‐template text
+        text = tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        # 3) Quick length check on token IDs
+        tokens = tokenizer(text, add_special_tokens=False)["input_ids"]
+        if len(tokens) <= max_len:
+            selected.append({"text": text})
+
+            # 4) Stop as soon as we have enough
+            if len(selected) >= n_samples:
+                break
+
+    # 5) Shuffle and build a regular Dataset
+    random.Random(seed).shuffle(selected)
+    return Dataset.from_list(selected)
+
+# --- usage in your script ---
+dataset = load_fitting_samples(
+    DATASET_ID,
+    tokenizer=tokenizer,
+    max_len=MAX_LEN,
+    n_samples=N_SAMPLES,
+    seed=3407,
+)
+
 new_dataset = dataset.train_test_split(test_size=0.2)
 
 # Configure training arguments