Create a Reasoning-Focused LLM: A Practical Guide to Streaming, Curating, and Fine-Tuning the SupraLabs Reasoning Corpus

In this tutorial, we build an end-to-end workflow for creating a reasoning-focused language model using the SupraLabs reasoning corpus (v1, containing 4K–5M samples). We stream a representative subset directly from the Hugging Face Hub, analyze its source distribution, token-length patterns, task composition, and reasoning-to-answer ratios, and apply quality filters to remove unsuitable examples. We then transform the retained samples into a chat-based supervised fine-tuning format with explicit reasoning tags and adapt SmolLM2-135M-Instruct using LoRA via TRL's SFTTrainer. This complete Google Colab pipeline covers scalable data access, exploratory analysis, dataset curation, parameter-efficient fine-tuning, structured inference, and Parquet export—turning a large multi-model reasoning corpus into a compact, task-oriented model.

1. Setting Up the Environment

We begin by installing the necessary libraries and configuring the runtime. The following commands install datasets, transformers, trl, peft, accelerate, bitsandbytes, matplotlib, and pandas. We also uninstall torchao to avoid potential conflicts with the training libraries.

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",
])

We then import all required modules, suppress warnings, set random seeds for reproducibility, and detect the compute device (CUDA or CPU). The dataset identifier and sample size are defined for the streaming step.

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 = 8000

2. Streaming the Dataset

The dataset is streamed directly from the Hugging Face Hub using the datasets library. We load a representative subset of 8,000 samples to explore its structure without downloading the entire corpus.

print(f"\nStreaming {DATASET_ID} ...")
stream = load_dataset(
    DATASET_ID,
    split="train",
    streaming=True
)
# Take a sample of SAMPLE_SIZE records
sample = list(itertools.islice(stream, SAMPLE_SIZE))
df = pd.DataFrame(sample)
print(f"Loaded {len(df)} samples.")

3. Exploratory Data Analysis

We inspect the dataset's column names, data types, and sample entries to understand its structure. Then we analyze:

• Source distribution: Which model or origin contributed the most samples?
• Token-length patterns: Are there outliers or typical lengths for reasoning and answer fields?
• Task composition: What types of reasoning tasks are present (e.g., math, logic, science)?
• Reasoning-to-answer ratios: How verbose is the reasoning relative to the answer?

We use pandas and matplotlib to generate summary statistics and visualizations.

print("Column names:", df.columns.tolist())
print(df.head())

# Source distribution
source_counts = df['source'].value_counts()
print(source_counts)

# Token-length analysis (assuming 'reasoning' and 'answer' columns)
if 'reasoning' in df.columns and 'answer' in df.columns:
    df['reasoning_len'] = df['reasoning'].apply(lambda x: len(str(x).split()))
    df['answer_len'] = df['answer'].apply(lambda x: len(str(x).split()))
    print(df[['reasoning_len', 'answer_len']].describe())
    df[['reasoning_len', 'answer_len']].hist(bins=50)
    plt.show()

4. Quality Filtering

To prepare the data for fine-tuning, we apply a series of filters to remove low-quality or unsuitable examples:

• Minimum and maximum token counts for reasoning and answer fields (e.g., reasoning 20–500 tokens, answer 1–200 tokens).
• Remove samples with empty or placeholder responses.
• Deduplicate identical reasoning chains if needed.
• Filter by source if some sources are known to be noisy.

We implement these rules as a function and apply them to the dataframe.

def filter_quality(df):
    # Example filters
    df = df[(df['reasoning_len'] >= 20) & (df['reasoning_len'] <= 500)]
    df = df[(df['answer_len'] >= 1) & (df['answer_len'] <= 200)]
    df = df[df['answer'].notna() & (df['answer'].str.strip() != '')]
    df = df.drop_duplicates(subset=['reasoning'])
    return df

df_clean = filter_quality(df)
print(f"Original: {len(df)}, After filtering: {len(df_clean)}")

5. Transforming to Chat Format with `` Tags

We convert the clean samples into a chat-based supervised fine-tuning format. For each example, we create a conversation with a user and assistant message, where the assistant message contains a reasoning block followed by the final answer. This structure helps the model learn to reason explicitly.

def format_chat(sample):
    prompt = sample['question']  # adjust based on actual column names
    reasoning = sample['reasoning']
    answer = sample['answer']
    return {
        "messages": [
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": f"<think>{reasoning}</think>\\n{answer}"}
        ]
    }

df_chat = df_clean.apply(lambda x: format_chat(x), axis=1)
# Convert to Hugging Face Dataset
from datasets import Dataset
dataset = Dataset.from_pandas(pd.DataFrame(df_chat.tolist()))

6. Fine-Tuning with LoRA and TRL

We load the base model (SmolLM2-135M-Instruct) and tokenizer with bfloat16 precision. We then configure LoRA (Low-Rank Adaptation) for parameter-efficient fine-tuning, targeting attention and feed-forward layers. Using TRL's SFTTrainer, we set training arguments and train the model for a few epochs on the curated dataset.

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from peft import LoraConfig
from trl import SFTTrainer

model_name = "HuggingFaceTB/SmolLM2-135M-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

training_args = TrainingArguments(
    output_dir="./smol_lora",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=2,
    num_train_epochs=3,
    learning_rate=2e-4,
    fp16=(DEVICE == "cuda"),
    save_total_limit=2,
    logging_steps=50,
    report_to="none"
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    peft_config=lora_config,
)

trainer.train()

7. Structured Inference and Export

After training, we run structured inference on a few test prompts to verify the model's reasoning behavior. The model should produce blocks followed by answers. We then export the adapted model and the curated dataset for future use.

# Inference example
prompts = ["What is the capital of France?", "Solve: 15 - 7 = ?"]
for prompt in prompts:
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    outputs = model.generate(**inputs, max_new_tokens=100)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# Save the adapted model
model.save_pretrained("./smol_lora_final")
tokenizer.save_pretrained("./smol_lora_final")

# Save the curated dataset as Parquet
df_clean.to_parquet("curated_reasoning_corpus.parquet")

Conclusion

This guide demonstrates a complete, production-ready pipeline for creating a reasoning-focused LLM. By streaming a large corpus efficiently, analyzing its characteristics, applying quality filters, and fine-tuning a compact model with LoRA, we produce a specialized model that can reason effectively in a structured format. The included Google Colab implementation is fully reproducible and can be adapted to other datasets and base models. With these tools, you can turn any sizable reasoning corpus into a nimble, focused language model suitable for deployment in resource-constrained environments.

via MarkTechPost