Создание LLM, ориентированной на рассуждения: практическое руководство по потоковой загрузке, отбору и дообучению корпуса рассуждений SupraLabs
В этом руководстве мы создадим сквозной рабочий процесс для работы с корпусом рассуждений SupraLabs. Мы загрузим репрезентативное подмножество напрямую из Hugging Face Hub в потоковом режиме, изучим распределение источников, закономерности длины последовательностей в токенах, состав задач и соотношение рассуждений и ответов, а затем применим ряд фильтров качества для удаления неподходящих примеров обучения. Мы преобразуем сохранённые образцы в формат контролируемого дообучения на основе чатов с явными тегами рассуждений <think> и используем их для адаптации SmolLM2-135M-Instruct с помощью LoRA через SFTTrainer из TRL. Объединив масштабируемый доступ к данным, исследовательский анализ, курирование датасета, эффективное дообучение параметров, структурированный вывод и экспорт в Parquet, мы создадим полноценный конвейер Google Colab для преобразования большого корпуса рассуждений на основе множества моделей в компактную языковую модель, ориентированную на рассуждения.
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]} ...")
Мы настроим среду Colab, установим необходимые библиотеки машинного обучения и удалим несовместимый пакет torchao. Мы определим доступное вычислительное устройство, подключимся к корпусу рассуждений SupraLabs через потоковую загрузку из Hugging Face и избежим скачивания всего датасета целиком. Мы перемешаем записи в потоке, материализуем репрезентативную выборку и изучим структуру и содержимое одной строки-примера.
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())
Мы преобразуем выбранный датасет в DataFrame pandas и анализируем распределение исходных репозиториев и длины последовательностей в токенах. Мы подсчитываем количество символов в рассуждениях и ответах, измеряем соотношение рассуждений и ответа и визуализируем взаимосвязи в датасете. Мы также применяем простые эвристические правила, чтобы классифицировать каждую запись как задачу по программированию, математике, медицине, задачу с вариантами ответа или общую задачу.
def filter_length(row, min_tok=200, max_tok=3000):
"""Keep samples within a training-friendly token budget."""
return min_tok <= row["tok_len"] <= max_tok
def filter_degenerate(row):
"""Drop empty/near-empty thoughts or answers."""
return len(row["thought_trace"]) > 100 and len(row["assistant"]) > 20
def filter_repetition(row, max_line_repeat=0.30):
"""Drop traces where one line repeats too often (looping models)."""
lines = [l.strip() for l in row["thought_trace"].split("\n") if l.strip()]
if len(lines) < 5:
return True
most_common = Counter(lines).most_common(1)[0][1]
return (most_common / len(lines)) <= max_line_repeat
def filter_reason_ratio(row, lo=0.15, hi=0.97):
"""Keep samples that actually reason but don't ONLY reason."""
t, a = len(row["thought_trace"]), len(row["assistant"])
r = t / (t + a + 1)
return lo <= r <= hi
n0 = len(ds)
ds_f = ds.filter(filter_length)
ds_f = ds_f.filter(filter_degenerate)
ds_f = ds_f.filter(filter_repetition)
ds_f = ds_f.filter(filter_reason_ratio)
print(f"\nFiltering: {n0:,} -> {len(ds_f):,} rows "
f"({100 * len(ds_f) / n0:.1f}% retained)")
MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
SYSTEM_PROMPT = (
"You are a careful reasoning assistant. Think step by step inside "
"<think>...</think> tags, then give your final answer."
)
def to_chat(row):
return {
"messages": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": row["user"]},
{"role": "assistant",
"content": f"<think>\n{row['thought_trace']}\n</think>\n\n{row['assistant']}"},
]
}
train_ds = ds_f.map(to_chat, remove_columns=ds_f.column_names)
train_ds = train_ds.shuffle(seed=42)
N_TRAIN, N_EVAL = 1_500, 100
eval_ds = train_ds.select(range(N_TRAIN, min(N_TRAIN + N_EVAL, len(train_ds))))
train_ds = train_ds.select(range(min(N_TRAIN, len(train_ds))))
print(f"\nTrain: {len(train_ds):,} | Eval: {len(eval_ds):,}")
print("\nRendered training sample (truncated):")
print(tokenizer.apply_chat_template(train_ds[0]["messages"], tokenize=False)[:800])
Мы создадим конвейер фильтрации качества, который удаляет примеры с неподходящей длиной в токенах, неполными ответами, чрезмерными повторами или несбалансированным содержанием рассуждений. Мы загрузим токенизатор SmolLM2 и преобразуем каждую сохранённую запись в структурированный диалог, содержащий системный запрос, сообщение пользователя и ответ ассистента с добавленными рассуждениями. Затем мы перемешаем отформатированные данные, создадим обучающую и оценочную выборки и изучим итоговый шаблон чата, используемый для контролируемого дообучения.
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig
try:
import peft.import_utils as _piu
import peft.tuners.lora.torchao as _plt
_piu.is_torchao_available = lambda: False
_plt.is_torchao_available = lambda: False
except Exception:
pass
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
).to(DEVICE)
peft_config = LoraConfig(
r=16,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
sft_config = SFTConfig(
output_dir="smollm2-reasoning-demo",
max_length=2048,
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
num_train_epochs=1,
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_steps=10,
logging_steps=10,
eval_strategy="steps",
eval_steps=50,
save_strategy="no",
bf16=(DEVICE == "cuda"),
gradient_checkpointing=True,
report_to="none",
)
trainer = SFTTrainer(
model=model,
args=sft_config,
train_dataset=train_ds,
eval_dataset=eval_ds,
peft_config=peft_config,
processing_class=tokenizer,
)
print("\nStarting fine-tune (≈10–20 min on a T4 with these settings)...")
trainer.train()
print("Done. Final eval loss:", trainer.evaluate().get("eval_loss"))
Мы загрузим языковую модель SmolLM2 для каузального моделирования и настроим адаптеры LoRA для эффективного обучения параметров. Мы зададим параметры оптимизации, формирования пакетов, оценки, точности и контрольных точек градиентов с помощью SFTConfig из TRL. Мы инициализируем SFTTrainer, дообучим модель на отобранных диалогах с рассуждениями и оценим итоговые показатели обучения.
def generate(question, max_new_tokens=512, temperature=0.7):
msgs = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": question},
]
prompt = tokenizer.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=True
)
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
with torch.no_grad():
out = trainer.model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=0.9,
do_sample=True,
pad_token_id=tokenizer.pad_token_id,
)
text = tokenizer.decode(out[0][inputs["input_ids"].shape[1]:],
skip_special_tokens=True)
m = re.search(r"<think>(.*?)</think>(.*)", text, re.DOTALL)
if m:
print("─" * 60, "\nTHINKING:\n", m.group(1).strip()[:1500])
print("─" * 60, "\nANSWER:\n", m.group(2).strip())
else:
print(text)
print("\n\n### TEST 1: logic puzzle")
generate("If all bloops are razzies and all razzies are lazzies, "
"are all bloops definitely lazzies? Explain briefly.")
train_ds.to_parquet("reasoning_subset_train.parquet")
eval_ds.to_parquet("reasoning_subset_eval.parquet")
print("\nSaved: reasoning_subset_train.parquet / reasoning_subset_eval.parquet")
Мы создадим функцию для вывода, которая форматирует новые вопросы с тем же системным запросом и генерирует ответы дообученной модели. Мы отделим сгенерированный раздел <think> от итогового ответа и протестируем модель на задачах по логике и арифметике. В завершение мы экспортируем обработанные обучающую и оценочную выборки в файлы Parquet для повторного использования в более крупных экспериментах.
В заключение мы разработали практический конвейер, связывающий исследование крупномасштабных данных для рассуждений с обучением небольшой языковой модели. Мы эффективно загрузили корпус в потоковом режиме, проанализировали его внутренний состав, отфильтровали примеры по критериям длины в токенах, повторов, полноты и сбалансированности рассуждений, а затем преобразовали полученные данные в единую структуру диалогов для обучения. После этого мы дообучили SmolLM2 с помощью LoRA, оценили адаптированную модель, изучили сгенерированные ею рассуждения и ответы и экспортировали курированные датасеты для будущих экспериментов. Этот рабочий процесс создаёт повторно используемую основу для смешивания данных с учётом источников, обучения по учебному плану, использования более крупных моделей-учеников, обучения с расширенным контекстом и разработки производственных моделей рассуждений без необходимости хранить весь датасет в памяти Colab.
Ознакомьтесь с ПОЛНЫМ КОДОМ здесь. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему сабреддиту по машинному обучению с более чем 150 тыс. участников и подписаться на нашу рассылку. Постойте! Вы есть в Telegram? теперь вы также можете присоединиться к нам в Telegram.
Хотите стать нашим партнёром для продвижения своего репозитория GitHub, страницы Hugging Face, выпуска продукта, вебинара и т. д.? Свяжитесь с нами
Переведено автоматически с английского. Оригинал статьи — по ссылке ниже.