The SupraLabs reasoning corpus is now available for streaming directly from the Hugging Face Hub. This dataset contains between 4,000 and 5,000,000 examples of model reasoning traces, structured to support the fine-tuning of smaller language models. The following workflow demonstrates how to stream a subset, inspect the data distribution, apply quality filters, and fine-tune the SmolLM2-135M-Instruct model using LoRA.
Setup and data access
The process begins by configuring the Google Colab environment. Required libraries, including datasets, transformers, and TRL, are installed. The incompatible torchao package is removed to prevent errors. The code detects the available compute device, defaulting to CPU if a GPU is not present.
Accessing the full dataset is unnecessary for this demonstration. Instead, the workflow streams the data from the repository identified as SupraLabs/reasoning-corpus-4K-5M-v1. A sample of 8,000 records is shuffled and materialised into a local dataset object. This approach saves memory and time compared to downloading the entire file.
import subprocess, sys
def pip_install(pkgs):
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *pkgs])
subprocess.call([sys.executable, "-m", "pip", "uninstall", "-y", "-q", "torchao"])
pip_install([
"datasets>=3.0.0",
"transformers>=4.46.0",
"trl>=0.12.0",
"peft>=0.13.0",
"accelerate>=1.0.0",
"bitsandbytes",
"matplotlib",
"pandas",
])
import os, re, json, math, random, itertools, warnings
import pandas as pd
import matplotlib.pyplot as plt
import torch
from collections import Counter
from datasets import load_dataset, Dataset
warnings.filterwarnings("ignore")
random.seed(42)
torch.manual_seed(42)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {DEVICE}")
if DEVICE == "cuda":
print(f"GPU: {torch.cuda.get_device_name(0)}")
DATASET_ID = "SupraLabs/reasoning-corpus-4K-5M-v1"
SAMPLE_SIZE = 8_000
print(f"\nStreaming {DATASET_ID} ...")
stream = load_dataset(DATASET_ID, split="train", streaming=True)
stream = stream.shuffle(seed=42, buffer_size=30_000)
rows = list(itertools.islice(stream, SAMPLE_SIZE))
ds = Dataset.from_list(rows)
print(f"Materialized sample: {len(ds):,} rows")
print(f"Columns: {ds.column_names}")
ex = ds[0]
print("\n" + "=" * 70)
print("EXAMPLE ROW")
print("=" * 70)
print(f"repo_id : {ex['repo_id']}")
print(f"tok_len : {ex['tok_len']}")
print(f"user : {ex['user'][:300]} ...")
print(f"thought_trace : {ex['thought_trace'][:300]} ...")
print(f"assistant : {ex['assistant'][:300]} ...")
Exploratory analysis
The sampled data is converted into a pandas DataFrame for analysis. The distribution of source repositories and token lengths is examined to understand the dataset composition. The code calculates the number of characters in the reasoning trace and the final answer for each record.
A reasoning ratio is computed by dividing the character count of the thought trace by the sum of the thought trace and answer lengths. Visualisations display the token length distribution, the top twelve source repositories, the reasoning ratio histogram, and a scatter plot correlating token length with the reasoning ratio.
Heuristic rules classify each record into categories such as code, mathematics, medical, multiple-choice logic, or general tasks. These rules scan the user prompt and assistant response for specific keywords and patterns.
df = ds.to_pandas()
print("\nTop 15 source repos in sample:")
src_counts = df["repo_id"].value_counts()
print(src_counts.head(15).to_string())
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
axes[0, 0].hist(df["tok_len"], bins=60, color="#4C72B0", edgecolor="white")
axes[0, 0].set_title("Token length distribution")
axes[0, 0].set_xlabel("tok_len"); axes[0, 0].set_ylabel("rows")
src_counts.head(12).plot(kind="barh", ax=axes[0, 1], color="#55A868")
axes[0, 1].invert_yaxis()
axes[0, 1].set_title("Top-12 source repos (sample)")
df["think_chars"] = df["thought_trace"].str.len()
df["answer_chars"] = df["assistant"].str.len()
df["reason_ratio"] = df["think_chars"] / (df["think_chars"] + df["answer_chars"] + 1)
axes[1, 0].hist(df["reason_ratio"], bins=50, color="#C44E52", edgecolor="white")
axes[1, 0].set_title("Reasoning ratio (think / (think + answer))")
axes[1, 0].set_xlabel("ratio")
axes[1, 1].scatter(df["tok_len"], df["reason_ratio"], s=4, alpha=0.25, color="#8172B2")
axes[1, 1].set_title("tok_len vs reasoning ratio")
axes[1, 1].set_xlabel("tok_len"); axes[1, 1].set_ylabel("ratio")
plt.tight_layout()
plt.show()
print("\nSummary stats:")
print(df[["tok_len", "think_chars", "answer_chars", "reason_ratio"]]
.describe().round(2).to_string())
def tag_task(row):
u = row["user"].lower()
a = row["assistant"]
if "```" in a or re.search(r"\b(def |class |import |function|#include)", a):
return "code"
if re.search(r"(prove|equation|integral|theorem|\\frac|\\int|solve for)", u):
return "math"
if re.search(r"\b(patient|diagnosis|symptom|treatment|clinical)\b", u):
return "medical"
if re.search(r"\b(which of the following|options?:|\(a\)|\(b\))", u):
return "mcq/logic"
return "general"
df["task"] = df.apply(tag_task, axis=1)
print("\nHeuristic task mix:")
print(df["task"].value_counts(normalize=True).round(3).to_string())
Quality filtering and formatting
The dataset undergoes a filtering process to remove unsuitable training examples. Four specific filters are applied sequentially. First, samples are kept only if their token length falls between 200 and 3,000 tokens. Second, records with empty or near-empty thoughts or answers are dropped. Third, traces where a single line repeats more than 30% of the time are removed to avoid looping models. Finally, samples are retained only if the reasoning ratio is between 0.15 and 0.97, ensuring the data contains both reasoning and final answers.
The remaining records are formatted into a chat-based structure. A system prompt instructs the model to think step by step inside <think> tags. The user prompt comes from the original user field, and the assistant response combines the thought_trace and assistant fields.
The formatted dataset is shuffled and split into training and evaluation subsets. The training set contains 1,500 examples, while the evaluation set holds 100 examples. The final chat template is rendered to verify the structure before fine-tuning.
def filter_length(rowSource Read original →Related articles
