Introduction to Model Distillation: Efficient Knowledge Transfer for AI Applications - Part 1¶
Pre requisites¶
- Nebius API key. Sign up for free at Token Factory
Introduction¶
Model distillation is a powerful technique in machine learning where a compact "student" model learns to replicate the behavior of a larger, more complex "teacher" model within the given task. By transferring knowledge from the teacher to the student, distillation enables lightweight models to achieve comparable performance within the task, while being dramatically faster and cheaper to deploy (and, consequently, cheaper to run inference with).
The benefits are compelling:
- Latency improvement: Smaller models perform much faster, which makes them ideal for real-time applications like agentic scenarios or other tasks where an immediate response is required.
- Cost reduction: Smaller models require less compute for inference, and hence are available at cheaper rates. Furthermore, the fine-tuned model removes the need in the long detailed prompts to ensure a specific format of the output data, which also reduces the price due to the lower tokens consumption.
In this tutorial, we demonstrate how to perform distillation using Nebius Token Factory to create a grammar-correcting model. We will:
- Generate high-quality training data via batched LLM generation using the recently released Qwen3-235B-A22B.
- Fine-tune a Qwen3-4B non-reasoning student model using LoRA adapters
- Deploy, evaluate and compare the distilled model with a 3.5x times larger model of this family, Qwen3-14B, using the most powerful open-source LLM to date, DeepSeek-R1, as evaluator.
By leveraging Nebius Token Factory’s batched generation, fine-tuning API, optimized inference and zero-click model deployment, we streamline the entire workflow—proving that large capabilities can indeed come in small packages. Let’s dive in!
Before we start, please note three things.
First, the procedure we employ differs from traditional distillation where the student model is trained on teacher's internal representations - instead, we will simply train the student model on the completions of the teacher model.
Second, the goal of this blog post is not to maximize the quality of the model on the given task but rather exhibit how to correctly perform distillation and why it matters. Hence, we will not focus on task-specific quality-optimization tricks or play around with the data. However, we will include all best practices for distillation so that your distilled model goes beyond your expectations!
Finally, due to non-deterministic parameters recommended by Qwen3 authors to run Qwen3, your results may slightly differ from the ones you see here if you relaunch the code. But no need to worry - we made sure the quality of the fine-tuned model is guaranteed to stay within the confidence interval of the baseline model!
1 - Prerequisites¶
Create an .env file with NEBIUS_API_KEY as follows
NEBIUS_API_KEY=your_api_key_goes_here
2 - Dependencies¶
Let's start with importing the necessary packages.
import os
from dotenv import load_dotenv
from typing import Sequence
from openai import Client
from datasets import load_dataset, Dataset, concatenate_datasets
from tqdm import tqdm
import pandas as pd
import json
import numpy as np
import time
import requests
import re
3 - Load Configuration¶
from dotenv import load_dotenv
load_dotenv()
True
4 - Initialize Client¶
The cell below creates the OpenAI-like Client to work with Nebius Token Factory and defines necessary variables.
DATASETS_CACHE_DIR = 'cache'
BASE_URL = "https://api.tokenfactory.nebius.com"
client = Client(
base_url=f'{BASE_URL}/v1',
api_key=os.getenv('NEBIUS_API_KEY')
)
5 - Load Input Dataset¶
In this tutorial, we want to demonstrate how, given only a dataset of input texts, train a small model by leveraging the most powerful LLMs to generate the desired outputs.
We will take a C4-200M dataset [1] for these purposes, which is intended for GEC models pre-training. Its outputs are anyway unsuitable for direct fine-tuning of a GEC model because it contains many errors, for example:
- Input: review narrow river as if air surf ...
- Output: air washer review boneco w200 air washer winix air washer review..
We will use its inputs and generate correct outputs using one of state-of-the-art LLMs - a recently released Qwen3-235B-A22B [2]. With a proper prompt tuning, we can urge the model output the data in easy-to-reuse format so that we can create the dataset to fine-tune our target small model -- Qwen3-4B [2].
Tens of thousands observations is generally enough to improve the quality of the model. Let's take a subset of 25k observations, process it by removing too short and too long sentences (this will leave us at 22k), and split into train and validation subsets for fine-tuning (21k & 1k).
input_dataset = load_dataset('Aktsvigun/c4_200m_25k', split='train', cache_dir=DATASETS_CACHE_DIR)
input_dataset
Dataset({
features: ['input'],
num_rows: 25000
})
Let's examine a random instance from the dataset.
input_dataset[2025]
{'input': 'Are you dissapointed on DNF or upsng race ettiraces?'}
C4-200M dataset is intended to contain sentences. Sentences below 3 or above 40 words are definite outliers, which most likely contain some garbage inputs. Let's filter out such input texts.
input_dataset = input_dataset.filter(lambda x: 40 > len(x['input'].split()) > 3)
input_dataset
Dataset({
features: ['input'],
num_rows: 22114
})
6 - Batch Inference¶
Heads up: Running this part will cost ~$4.9.
You can use normal synchronous generations with Qwen3-235B-A22B to generate outputs for the dataset. However, if you are not in a last-minute rush, usage of batch inference is recommended. It provides as much as 2x cheaper rates, and is guaranteed to finish within 24 hours. In most cases, it takes a few hours or even minutes, again, depending on the size of the dataset.
Let's see how to use the batched generation to annotate our input dataset.
First, we need a carefully designed prompt to have the data generated in the desired format. Desired format here is untouched input sentence if it is already gramatically correct, or its corrected version, otherwise.
To urge the model follow the desired format for generation (without adding introduction like "Here is the corrected text" or further explanations), we will leverage few-shot learning examples. We provide one example per gramatically correct and incorrect input texts in our prompt.
system_prompt_distillation = """
Act as an experienced English proofreader. Please check the grammar of the user's text. If the text contains errors or misprints, print the corrected text. Otherwise, print the text as it is, otherwise. Check only the grammar of the text. Don't print anything else.
Examples for few-shot learning:
Example 1 (the text contains errors):
User: In fact who let me know abut this program was him.
Assistant: In fact, he was the one who let me know about this program.
Example 2 (the text does not contain errors):
User: On the other hand, it's very efficient computationally as it only requires one forward pass through the model per example.
Assistant: On the other hand, it's very efficient computationally as it only requires one forward pass through the model per example.
""".strip()
6.1 - Save input data in JSONL format¶
Let's format the dataset and save it as a .jsonl file.
!mkdir data
max_tokens = 4096
with open('data/batch_input.jsonl', 'w') as f:
for i, inst in enumerate(input_dataset, 1):
dict_to_write = {
"custom_id": f"request-{i}",
"method": "POST",
"url": "/v1/chat/completions",
"body": {
"model": "Qwen/Qwen3-235B-A22B-Instruct-2507",
"messages": [
{"role": "system", "content": system_prompt_distillation},
{"role": "user", "content": inst["input"]}
],
"max_tokens": max_tokens
}
}
json.dump(dict_to_write, f, ensure_ascii=False)
f.write('\n')
mkdir: cannot create directory ‘data’: File exists
6.2 - Upload our input dataset to Nebius Token Factory.¶
batch_input_file = client.files.create(
file=open("data/batch_input.jsonl", "rb"),
purpose="batch"
)
batch_input_file
FileObject(id='file-386da9e1-b2c8-42aa-a799-363b61cff93c', bytes=24630543, created_at=1753305500, filename='batch_input.jsonl', object='file', purpose='batch', status=None, expires_at=None, status_details=None)
6.3 - Launch the batch inference job¶
Now that all perliminary steps are done, use the uploaded dataset to create the batched generation job. The code below launches the batched generation.
batch_input_file_id = batch_input_file.id
batch = client.batches.create(
input_file_id=batch_input_file_id,
endpoint="/v1/chat/completions",
completion_window="24h",
metadata={
"description": "Distillation of Qwen/Qwen3-235B-A22B-Instruct-2507 for GEC"
}
)
batch
Batch(id='batch_624bbde2-58f6-42e9-adf4-096ca8a98569', completion_window='24h', created_at=1753305501, endpoint='/v1/chat/completions', input_file_id='file-386da9e1-b2c8-42aa-a799-363b61cff93c', object='batch', status='validating', cancelled_at=None, cancelling_at=None, completed_at=None, error_file_id=None, errors=None, expired_at=None, expires_at=None, failed_at=None, finalizing_at=None, in_progress_at=None, metadata={'description': 'Distillation of Qwen/Qwen3-235B-A22B for GEC'}, output_file_id=None, request_counts=BatchRequestCounts(completed=None, failed=None, total=None))
6.4 - Wait for batch inference job to finish¶
It will now take some time to complete your job. The total time highly depends on the workload of the model. In our case, it finished within 1 hour.
You can periodically monitor the status of your job. When the job is completed, status will be equal to 'done'. The cell below will update its status every minute and stop running once the job is finished.
%%time
import time
start_time = time.time()
update_num_seconds = 60
active_statuses = ["validating", "validated", "running"]
print (f"Batch {batch.id} created, waiting for completion...")
while batch.status in active_statuses:
time.sleep(update_num_seconds)
# Retrieve the batch state
batch = client.batches.retrieve(batch.id)
elapsed = time.time() - start_time
print(f"Elapsed: {int(elapsed)}s ({elapsed/60:.1f} min) : current status: {batch.status}")
Batch batch_624bbde2-58f6-42e9-adf4-096ca8a98569 created, waiting for completion... Elapsed: 60s (1.0 min) : current status: running Elapsed: 121s (2.0 min) : current status: running Elapsed: 182s (3.0 min) : current status: running Elapsed: 242s (4.0 min) : current status: running Elapsed: 303s (5.1 min) : current status: running Elapsed: 364s (6.1 min) : current status: running Elapsed: 425s (7.1 min) : current status: running Elapsed: 486s (8.1 min) : current status: running Elapsed: 547s (9.1 min) : current status: running Elapsed: 607s (10.1 min) : current status: running Elapsed: 668s (11.1 min) : current status: running Elapsed: 729s (12.2 min) : current status: running Elapsed: 789s (13.2 min) : current status: running Elapsed: 850s (14.2 min) : current status: running Elapsed: 911s (15.2 min) : current status: running Elapsed: 972s (16.2 min) : current status: running Elapsed: 1033s (17.2 min) : current status: running Elapsed: 1093s (18.2 min) : current status: running Elapsed: 1154s (19.2 min) : current status: running Elapsed: 1215s (20.3 min) : current status: running Elapsed: 1276s (21.3 min) : current status: running Elapsed: 1337s (22.3 min) : current status: running Elapsed: 1398s (23.3 min) : current status: running Elapsed: 1459s (24.3 min) : current status: running Elapsed: 1519s (25.3 min) : current status: running Elapsed: 1580s (26.3 min) : current status: running Elapsed: 1641s (27.4 min) : current status: running Elapsed: 1701s (28.4 min) : current status: running Elapsed: 1762s (29.4 min) : current status: running Elapsed: 1823s (30.4 min) : current status: running Elapsed: 1884s (31.4 min) : current status: running Elapsed: 1945s (32.4 min) : current status: running Elapsed: 2005s (33.4 min) : current status: running Elapsed: 2066s (34.4 min) : current status: running Elapsed: 2127s (35.5 min) : current status: running Elapsed: 2188s (36.5 min) : current status: running Elapsed: 2249s (37.5 min) : current status: running Elapsed: 2310s (38.5 min) : current status: running Elapsed: 2371s (39.5 min) : current status: running Elapsed: 2432s (40.5 min) : current status: running Elapsed: 2493s (41.6 min) : current status: running Elapsed: 2554s (42.6 min) : current status: running Elapsed: 2614s (43.6 min) : current status: running Elapsed: 2675s (44.6 min) : current status: running Elapsed: 2736s (45.6 min) : current status: running Elapsed: 2798s (46.6 min) : current status: running Elapsed: 2859s (47.7 min) : current status: running Elapsed: 2920s (48.7 min) : current status: running Elapsed: 2980s (49.7 min) : current status: running Elapsed: 3041s (50.7 min) : current status: running Elapsed: 3102s (51.7 min) : current status: running Elapsed: 3163s (52.7 min) : current status: running Elapsed: 3223s (53.7 min) : current status: running Elapsed: 3284s (54.7 min) : current status: running Elapsed: 3345s (55.8 min) : current status: running Elapsed: 3406s (56.8 min) : current status: running Elapsed: 3467s (57.8 min) : current status: running Elapsed: 3528s (58.8 min) : current status: running Elapsed: 3588s (59.8 min) : current status: running Elapsed: 3649s (60.8 min) : current status: running Elapsed: 3710s (61.8 min) : current status: running Elapsed: 3771s (62.9 min) : current status: running Elapsed: 3831s (63.9 min) : current status: running Elapsed: 3892s (64.9 min) : current status: running Elapsed: 3953s (65.9 min) : current status: running Elapsed: 4014s (66.9 min) : current status: running Elapsed: 4075s (67.9 min) : current status: running Elapsed: 4135s (68.9 min) : current status: running Elapsed: 4196s (69.9 min) : current status: running Elapsed: 4257s (71.0 min) : current status: running Elapsed: 4318s (72.0 min) : current status: running Elapsed: 4378s (73.0 min) : current status: running Elapsed: 4439s (74.0 min) : current status: running Elapsed: 4500s (75.0 min) : current status: running Elapsed: 4561s (76.0 min) : current status: running Elapsed: 4622s (77.0 min) : current status: running Elapsed: 4683s (78.1 min) : current status: running Elapsed: 4743s (79.1 min) : current status: running Elapsed: 4804s (80.1 min) : current status: running Elapsed: 4865s (81.1 min) : current status: running Elapsed: 4926s (82.1 min) : current status: running Elapsed: 4987s (83.1 min) : current status: running Elapsed: 5047s (84.1 min) : current status: running Elapsed: 5108s (85.1 min) : current status: running Elapsed: 5169s (86.2 min) : current status: running Elapsed: 5230s (87.2 min) : current status: done CPU times: user 1.09 s, sys: 136 ms, total: 1.22 s Wall time: 1h 27min 10s
6.5 - Save Batch inference data locally¶
Our batch has been successfully processed. Let's save the generations to a file and examine the format in which it comes.
file_response = client.files.content(batch.output_file_id)
file_response.write_to_file('data/batch_output.jsonl')
# Display the first line from the output file
with open('data/batch_output.jsonl') as f:
for line in f.readlines():
output = json.loads(line)
break
print(json.dumps(output, indent=4))
print ("✅ Batch processing completed successfully!")
{
"id": "batch_req_fc038ea5-6112-4b0f-be71-e402f74a8105",
"custom_id": "request-1580",
"response": {
"id": "chatcmpl-c03d469ec6bd4c61a7a9a22157ce548b",
"choices": [
{
"finish_reason": "stop",
"index": 0,
"logprobs": null,
"message": {
"content": "<think>\nOkay, let's tackle this query. The user wants me to act as an experienced English proofreader. The task is to check the grammar of their text. If there are errors or misprints, I need to correct them. If not, just return the text as is. Important: only check grammar, don't add anything else.\n\nFirst, I'll read through the user's text carefully. The example they provided shows that the assistant should correct grammatical errors and misprints, like typos or wrong word choices, but leave the technical terms and product names intact. \n\nThe user's text is: \"Fluke Australia - Th7 Fluke 1587 and 1577 insulations Multimeters combine a digital multimeter in a single compact unitted unit with their heads which provides maximum verse for both troubleshooting and preventative maintenance.\"\n\nLooking at the first part: \"Th7 Fluke 1587\". \"Th7\" might be a model number or part of the product name. Since the example had \"Fluke 1587\" corrected but \"Fluke\" remained, maybe \"Th7\" is intentional. However, the hyphen after \"Australia\" could be correct if it's part of the company's branding. I should check if \"Th7\" is a typo. Maybe it's supposed to be \"The Fluke\"? But in the example, the company name was left as is. So probably \"Th7\" is correct here.\n\nNext, \"insulations Multimeters\". The word \"insulations\" seems off. It should likely be \"insulation\" as it's part of the product name, maybe \"Insulation Multimeters\". So \"insulation Multimeters\" would be correct.\n\nThen, \"combine a digital multimeter in a single compact unitted unit\". The phrase \"unitted unit\" is redundant. \"Unitted\" is probably a typo. The correct term might be \"united\", but that still doesn't make sense. More likely, it should be \"unit\" instead of \"unitted unit\". Wait, the original text says \"compact unitted unit\". \"Unitted\" isn't a standard word here. Looking at the example, \"unitted\" was corrected to \"integrated\". Wait, in the example, the user wrote \"unitted unit\" and the assistant corrected it to \"integrated\". Wait, no, in example 1, the correction was from \"unitted unit\" to \"integrated\". Wait, no, the example given was different. Let me check again. The user's example 1 had \"unit\" corrected to \"unit\". Wait, no. Wait, the example 1 user wrote: \"In fact who let me know abut this program was him.\" The assistant corrected it to \"In fact, he was the one who let me know about this program.\" So the assistant added commas, corrected \"abut\" to \"about\", restructured the sentence. So in the current user's text, \"unitted unit\" \u2013 \"unitted\" might be a typo for \"united\" or \"integrated\"? Or perhaps \"compact unit\" is intended, so \"unitted\" is an error. The correct phrase might be \"a single compact unit\". So \"unitted\" is a typo for \"unit\", but the original says \"unitted unit\". So maybe it should be \"single compact unit\" without \"unitted\". Alternatively, \"compact unit\" with \"unitted\" being redundant. Alternatively, \"integrated\" instead of \"unitted\". Since \"integrated\" is a common term in product descriptions, maybe that's the intended word. So \"integrated unit\".\n\nThen, \"with their heads which provides maximum verse\". \"Heads\" here might be a typo. If it's referring to the device's components, maybe \"head\" (singular) or another term like \"handles\" or \"probes\"? But \"heads\" could be correct if that's the technical term. However, \"provides\" should be \"provide\" because \"heads\" is plural. So \"which provide\". Also, \"maximum verse\" \u2013 \"verse\" doesn't make sense here. The correct word is likely \"versatility\", so \"maximum versatility\".\n\nPutting it all together:\n\nOriginal: \"Fluke Australia - Th7 Fluke 1587 and 1577 insulations Multimeters combine a digital multimeter in a single compact unitted unit with their heads which provides maximum verse for both troubleshooting and preventative maintenance.\"\n\nCorrections:\n- \"insulations Multimeters\" \u2192 \"insulation Multimeters\" (singular)\n- \"unitted unit\" \u2192 \"unit\" (removing \"unitted\")\n- \"heads which provides\" \u2192 \"heads, which provide\" (comma and plural verb)\n- \"verse\" \u2192 \"versatility\"\n\nSo the corrected sentence would be:\n\n\"Fluke Australia - Th7 Fluke 1587 and 1577 insulation Multimeters combine a digital multimeter in a single compact unit with their heads, which provide maximum versatility for both troubleshooting and preventative maintenance.\"\n\nWait, but in the original, \"Th7 Fluke 1587\" \u2013 maybe \"Th7\" is a typo for \"The\"? Like \"The Fluke 1587\"? Let me check the original again. The user wrote \"Th7 Fluke 1587\". If \"Th7\" is a model number, it's okay. But if it's a typo for \"The\", then it should be corrected. However, since the example assistant didn't correct the company name but fixed other errors, perhaps \"Th7\" is part of the product name. So leave it as is.\n\nAnother thing: \"preventative maintenance\" \u2013 \"preventative\" is a variant of \"preventive\". Both are correct, but \"preventive\" is more common. However, the user used \"preventative\", which is acceptable. So no correction needed there.\n\nAnother check: \"combine a digital multimeter in a single compact unit\" \u2013 does that make sense? Maybe \"combine a digital multimeter and...\" or \"combine a digital multimeter with...\". The original structure is a bit awkward. Wait, the original sentence structure is \"combine a digital multimeter in a single compact unit with their heads...\". Maybe it's supposed to be \"combine a digital multimeter and insulation tester in a single compact unit with their heads...\" but the original says \"insulation Multimeters\". Hmm. The phrase \"combine a digital multimeter in a single compact unit with their heads\" might need restructuring. Maybe it's supposed to be \"combine a digital multimeter and insulation tester in a single compact unit with their heads...\" but the original says \"insulation Multimeters combine...\". Maybe the intended structure is that the multimeters combine a digital multimeter and insulation tester in one unit. But the original text might have \"insulation Multimeters\" as a category. This part might be correct, but the grammar around it is off.\n\nSo focusing strictly on grammar and not technical accuracy, the main corrections are:\n\n1. \"insulations Multimeters\" \u2192 \"insulation Multimeters\" (singular)\n2. \"unitted unit\" \u2192 \"unit\" (remove \"unitted\")\n3. \"heads which provides\" \u2192 \"heads, which provide\" (comma and verb agreement)\n4. \"verse\" \u2192 \"versatility\"\n\nSo the corrected text would be:\n\nFluke Australia - Th7 Fluke 1587 and 1577 insulation Multimeters combine a digital multimeter in a single compact unit with their heads, which provide maximum versatility for both troubleshooting and preventative maintenance.\n\nI need to make sure that all grammatical errors are addressed. Let me check again.\n\n- \"insulations\" \u2192 \"insulation\" (correcting plural to singular for proper noun usage)\n- \"unitted unit\" \u2192 \"unit\" (removing \"unitted\" as it's a typo or redundancy)\n- \"heads which provides\" \u2192 \"heads, which provide\" (relative clause needs a comma, and verb agrees with plural \"heads\")\n- \"verse\" \u2192 \"versatility\" (correcting typo)\n\nYes, these changes address the grammatical issues. The rest of the text seems okay. \"Th7\" remains as it's part of the product name, and \"preventative\" is acceptable.\n</think>\n\nFluke Australia - Th7 Fluke 1587 and 1577 insulation Multimeters combine a digital multimeter in a single compact unit with their heads, which provide maximum versatility for both troubleshooting and preventative maintenance.",
"refusal": null,
"role": "assistant",
"annotations": null,
"audio": null,
"function_call": null,
"tool_calls": [],
"reasoning_content": null
},
"stop_reason": null
}
],
"created": 1753306953,
"model": "Qwen/Qwen3-235B-A22B",
"object": "chat.completion",
"service_tier": null,
"system_fingerprint": null,
"usage": {
"completion_tokens": 1777,
"prompt_tokens": 229,
"total_tokens": 2006,
"completion_tokens_details": null,
"prompt_tokens_details": null
},
"prompt_logprobs": null
},
"error": null
}
✅ Batch processing completed successfully!
7 - Data Cleanup¶
To get a model suitable for online application, let's query only the generations without the thinking part. Next, create a dataset that we'll afterwards merge with the input dataset.
Even though our dataset is not that large, let's create the Dataset object from file so that at no point we store the whole dataset in RAM - this will be a helpful example of how to deal with large datasets.
Since we want to exhibit a distillation for real-world usecases, we will only train the model on completions, discarding the thinking part. This ensures the responses are generated immediately, which is generally crucial for production applications. Hence, we extract the content after the </think> tag to save only the final corrected version.
There may also be cases where the model thought for so long that it didn't reach the final output. We will filter these cases by removing observations where the number of completion tokens coincides with maximum tokens we used for generation (4096).
output_save_path = 'data/batch_output_processed.jsonl'
prompt_tokens = 0
completion_tokens = 0
ids_to_filter = set()
with open(output_save_path, 'w') as f_out:
with open('data/batch_output.jsonl') as f_in:
for line in f_in.readlines():
output = json.loads(line)
output_text = output['response']['choices'][0]['message']['content']
output_without_thinking = output_text.split('</think>')[-1].strip()
output_id = int(output['custom_id'].split('-')[1])
# Check the generation was finished. We won't remove these instances at the moment:
# we will remove them once we concatenate the outputs with the input dataset
if output['response']['usage']['completion_tokens'] == max_tokens:
ids_to_filter.add(output_id)
json.dump({'output': output_without_thinking, 'id': output_id}, f_out, ensure_ascii=False)
f_out.write('\n')
# Calculate token statistics
prompt_tokens += output['response']['usage']['prompt_tokens']
completion_tokens += output['response']['usage']['completion_tokens']
8 - Examine the price of batch inference¶
Let's also calculate the price of the batched generation. We can take the price for input/output tokens of a model from the Token Factory home page. For Qwen/Qwen3-235B-A22B, it is \$0.2/\$0.6 for 1M input/output tokens. However, thanks to using batched generation, it costs twice as little with \$0.1/\$0.3 for 1M input/output tokens
price = (prompt_tokens * 0.1 + completion_tokens * 0.3) / 1_000_000
print(f'Batched generation price: ${price:.1f}')
Batched generation price: $4.9
9 - Merge and Clean the dataset¶
Now let's merge our dataset with outputs with dataset with inputs, remove the instance for which the generation has not been finished, and check that the merge didn't break anything.
output_dataset = Dataset.from_json(output_save_path, split="train")
output_dataset = output_dataset.sort('id')
output_dataset
Generating train split: 0 examples [00:00, ? examples/s]
Dataset({
features: ['output', 'id'],
num_rows: 22114
})
assert len(input_dataset) == len(output_dataset)
ft_dataset = concatenate_datasets([input_dataset, output_dataset], axis=1)
# Filter out unfinished generations
ft_dataset = ft_dataset.filter(lambda x: x['id'] not in ids_to_filter)
# Remove the `id` column, which is not useful anymore
ft_dataset = ft_dataset.remove_columns('id')
ft_dataset
Flattening the indices: 0%| | 0/22114 [00:00<?, ? examples/s]
Flattening the indices: 0%| | 0/22114 [00:00<?, ? examples/s]
Filter: 0%| | 0/22114 [00:00<?, ? examples/s]
Dataset({
features: ['input', 'output'],
num_rows: 22092
})
ft_dataset[42]
{'input': 'I think we need both 48bit & softprin in Libdrm.',
'output': 'I think we need both 48-bit and softprin in Libdrm.'}
10 - Save the final finetuning dataset¶
## Save final dataset in JSONL format
# ft_dataset.to_json('data/ft_dataset2.jsonl', orient="records", lines=True)
with open('data/ft_dataset.jsonl', 'w') as f:
for i, record in enumerate(ft_dataset):
json.dump(record, f, ensure_ascii=False)
f.write('\n')
print("✅ Final dataset saved to 'data/ft_dataset.jsonl'")
✅ Final dataset saved to 'data/ft_dataset.jsonl'
Our dataset for fine-tuning is created! We can now procede to fine-tuning - the most exciting part for most of AI developers :-).