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

from dotenv import load_dotenv
load_dotenv()

# read model name from file
model_name_file = 'model_name.out'
if os.path.exists(model_name_file):
    with open(model_name_file, 'r') as f:
        MODEL_NAME = f.read().strip()
else:
    # fallback to default model name if file does not exist
    raise Exception (f"Model name file '{model_name_file}' not found. Please run the fine-tuning notebook (2_training.ipynb) first.")

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')
)

system_prompt_fine_tuning = "Please correct the grammar in the user's text if necessary."

empty_reasoning_prefix = """
<think>

</think>

""".lstrip()

sample_text = "Nebius Token Factory is a comprehensive platform designed to simplify the integration of AI capabilities into applications."
text_with_errors = "Nebius Token Factory is comprehensive platform designed too simplify integration of AI capabiltes into applcations."

resp = client.chat.completions.create(
    model=MODEL_NAME,
    messages=[
        {'role': 'system', 'content': system_prompt_fine_tuning},
        {'role': 'user', 'content': text_with_errors}
    ],
    temperature=0.7,
    top_p=0.8,
    max_tokens=40
)
model_generation = resp.choices[0].message.content.split('</think>')[-1].strip()

if sample_text == model_generation:
    print('Model generation coincides with the original text!')
else:
    print('Model generation differs from the original text:', model_generation)

Model generation coincides with the original text!

eval_dataset = load_dataset('jhu-clsp/jfleg', cache_dir=DATASETS_CACHE_DIR)['test']
# Remove bad instances where no corrections are suggested or the sentence is empty
eval_dataset = eval_dataset.filter(lambda x: x['sentence'] != '' and not all(y == '' for y in x['corrections']))
print(eval_dataset[0])
eval_dataset

Generating validation split:   0%|          | 0/755 [00:00<?, ? examples/s]

Generating test split:   0%|          | 0/748 [00:00<?, ? examples/s]

Filter:   0%|          | 0/748 [00:00<?, ? examples/s]

{'sentence': 'New and new technology has been introduced to the society .', 'corrections': ['New technology has been introduced to society .', 'New technology has been introduced into the society .', 'Newer and newer technology has been introduced into society .', 'Newer and newer technology has been introduced to the society .']}

Dataset({
    features: ['sentence', 'corrections'],
    num_rows: 747
})

def fix_spacing(text):
    return re.sub(r'\s+([.,:;?!])', r'\1', text).strip()

eval_dataset = eval_dataset.map(lambda x: {"sentence": fix_spacing(x["sentence"])}, remove_columns="sentence")
eval_dataset = eval_dataset.map(lambda x: {"corrections": list(map(fix_spacing, x["corrections"]))}, remove_columns="corrections")
eval_dataset

Map:   0%|          | 0/747 [00:00<?, ? examples/s]

Map:   0%|          | 0/747 [00:00<?, ? examples/s]

Dataset({
    features: ['sentence', 'corrections'],
    num_rows: 747
})

%%time 

from concurrent.futures import ThreadPoolExecutor, as_completed

def _call_llm(input_text: str, system_prompt: str, **call_kwargs):
    return client.chat.completions.create(
        messages=[
            {'role': 'system', 'content': system_prompt},
            {'role': 'user', 'content': input_text},
            {'role': 'assistant', 'content': empty_reasoning_prefix}
        ],
        extra_body={"continue_final_message": True, "add_generation_prompt": False},
        **call_kwargs
    ).choices[0].message.content

def generate(
    input_text: str,
    model: str = MODEL_NAME,
    system_prompt: str = system_prompt_fine_tuning,
    max_tokens: int = 128,  # setting to 128 as default to make sure our generation won't terminate on long sentences
    top_p: float = 0.8,  # as suggested by Qwen3 authors
    temperature: float = 0.7,  # as suggested by Qwen3 authors
) -> str:
    # Use try-except construction in case we encounter an error 
    output = None
    while output is None:
        try:
            return _call_llm(
                input_text=input_text,
                system_prompt=system_prompt,
                model=model,
                max_tokens=max_tokens,
                top_p=top_p
            )
        except Exception as e:
            print(f'Error for input text {input_text}:\n{e}\nSleeping for 3 seconds and trying again...')
            time.sleep(3)
            continue
        
ft_model_generations = [None] * len(eval_dataset)
# Be careful with quota limits - you may need to reduce it
max_workers = 12
with ThreadPoolExecutor(max_workers=max_workers) as executor:
    futures = {executor.submit(generate, inst['sentence'], MODEL_NAME, system_prompt_fine_tuning): idx for idx, inst in enumerate(eval_dataset)}
    for future in tqdm(as_completed(futures), total=len(eval_dataset)):
        idx = futures[future]
        ft_model_generations[idx] = future.result()
ft_model_generations[:2]

100%|██████████| 747/747 [02:01<00:00,  6.17it/s]

['<think>\nOkay, let\'s see. The user wrote: "New and new technology has been introduced to the society." The user wants me to correct the grammar if necessary.\n\nFirst, I need to check if there are any grammatical errors. The phrase "New and new technology" seems repetitive. The word "new" is used twice. That\'s probably a redundancy. So maybe it should be "New technology" instead of "New and new technology." \n\nNext, the verb "has been introduced" is correct, but the subject "New and new technology" is singular. Wait, but "technology" is uncountable,',
 '<think>\nOkay, the user provided a sentence and wants me to correct any grammar issues. Let me read through it carefully.\n\nThe original sentence: "One possible outcome is that an environmentally-induced reduction in motorization levels in the richer countries will outweigh any rise in motorization levels in the poorer countries."\n\nFirst, I need to check for grammatical errors. Let\'s start with the structure. The sentence is complex, so I should ensure that the clauses are correctly connected. The main clause is "One possible outcome is that..." followed by a dependent clause. The dependent clause uses "that" correctly.\n\nLooking at the words: "environmentally']

%%time 

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()

bs_model_generations = [None] * len(eval_dataset)
# Reduce the parallelization to prevent rate limit errors
max_workers = 16
with ThreadPoolExecutor(max_workers=max_workers) as executor:  
    futures = {
        executor.submit(
            generate,
            inst['sentence'],
            'Qwen/Qwen3-30B-A3B-Instruct-2507',
            system_prompt_distillation,
        ): idx
        for idx, inst in enumerate(eval_dataset)
    }
    for future in tqdm(as_completed(futures), total=len(eval_dataset)):
        idx = futures[future]
        bs_model_generations[idx] = future.result()
bs_model_generations[:2]

100%|██████████| 747/747 [01:36<00:00,  7.77it/s]

CPU times: user 5.04 s, sys: 318 ms, total: 5.36 s
Wall time: 1min 36s

['<think>\nOkay, let\'s take a look at the user\'s sentence: "New and new technology has been introduced to the society." The user wants me to check the grammar. \n\nFirst, I notice "New and new technology." That repetition of "new" seems off. Maybe they meant "New and emerging technology" or "New and innovative technology." But the user might just be emphasizing the novelty, so maybe "new" is intentional. However, in standard grammar, repeating the same adjective isn\'t typically correct unless for stylistic reasons. But the user\'s example shows they correct repetition, so I should check that.\n\nNext,',
 '<think>\nOkay, let me take a look at the user\'s sentence. The original text is: "One possible outcome is that an environmentally-induced reduction in motorization levels in the richer countries will outweigh any rise in motorization levels in the poorer countries."\n\nFirst, I need to check for grammar errors. The sentence structure seems correct. The main clause is "One possible outcome is that..." followed by a subordinate clause. The use of "will outweigh" is appropriate here because it\'s a future prediction. \n\nLooking at the adjectives: "richer countries" and "poorer countries" are correctly comparative forms. The phrase "environment']

ft_results = {}
bs_results = {}
ft_ids_require_verification = []
bs_ids_require_verification = []
coinciding_ids = set()

for i in range(len(eval_dataset)):
    row = eval_dataset[i]
    sentence = row['sentence']
    corrections = row['corrections']
    # If generation is in corrections, it is considered valid
    if ft_model_generations[i] in corrections:
        ft_results[i] = 1
    # Else, if generation coincides with the original sentence, it is considered 
    # invalid  because the original sentence is not among the corrections
    # (otherwise it would meet the previous condition)
    elif ft_model_generations[i] == sentence:
        ft_results[i] = 0
    # Otherwise need to call an LLM for verification
    else:
        ft_ids_require_verification.append(i)

    # Same for the baseline model
    if bs_model_generations[i] in corrections:
        bs_results[i] = 1
    elif bs_model_generations[i] == sentence:
        bs_results[i] = 0
    # If the generation doesn't fall under the two conditions and coincides with the fine-tuned model generation,
    # we don't need to duplicate verification for the baseline model
    elif bs_model_generations[i] == ft_model_generations[i]:
        coinciding_ids.add(i)
    else:
        bs_ids_require_verification.append(i)
len(ft_ids_require_verification), len(bs_ids_require_verification)

CPU times: user 37.5 ms, sys: 2.03 ms, total: 39.5 ms
Wall time: 38.5 ms

(747, 747)

%%time 

evaluation_model = 'deepseek-ai/DeepSeek-R1-0528'

system_prompt_evaluate = """
Act as an experienced grammar checker.

You will be provided with:
1. A sentece that may contain errors
2. Its 4 corrections suggested by 4 language experts (corrections may coincide with the sentence if the expert believes the sentence is grammatically correct)
3. Its correction suggested by the proofreader

Please evaluate whether the correction written by a proofreader is valid (1) or not (0). \
The correction is considered invalid if it omits editing any of the errors or rectifies indeed gramatically correct pieces of text. \
Otherwise, it is considered valid.

Again, please output:
- "1", if the editing suggested by the proofreader is correct
- "0", if the editing suggested by the proofreader is incorrect

Only output the number ("0" / "1") and nothing else.
""".strip()

user_prompt_evaluate_template = """
Sentence:
{sentence}

Corrections:
{corrections}

Proofreader's correction:
{generation}
""".strip()

def evaluate_text(
    row: dict[str, str | list[str]], generation: str, system_prompt: str = system_prompt_evaluate
) -> str:
    formatted_corrections = '\n'.join(f"{i}. {correction}" for i, correction in enumerate(row['corrections']))    
    user_prompt = str(user_prompt_evaluate_template).format(
        sentence=row['sentence'],
        corrections=formatted_corrections,
        generation=generation
    )
    answer_with_reasoning = client.chat.completions.create(
        model=evaluation_model,
        messages=[
            {'role': 'system', 'content': system_prompt_evaluate},
            {'role': 'user', 'content': user_prompt},
        ],
        max_tokens=4096,  # to make sure the model finishes the reasoning and outputs an answer
        top_p=0.01,  # to make the results as deterministic as possible
    ).choices[0].message.content
    answer = answer_with_reasoning.split('</think>')[-1].strip()
    return convert_verdict_to_number(answer)

def convert_verdict_to_number(verdict: str) -> int:
    if verdict.isdigit():
        return int(verdict)
    # To make sure we don't overrate the performance of the model, if DeepSeek-R1 outputs smth else
    # than a number, consider the editing of the model invalid
    print(f"Cannot parse the answer {verdict}. Replacing with 0.")
    return 0

def evaluate(ids_require_verification: Sequence[int], model_generations: list[str]) -> str:
    data_require_verification = eval_dataset.select(ids_require_verification)
    model_labels = [None] * len(ids_require_verification)
    max_workers = 16
    with ThreadPoolExecutor(max_workers=max_workers) as executor:  # Reduce the parallelization by 8 times
        futures = {
            executor.submit(
                evaluate_text,
                data_require_verification[idx],
                generation
            ): idx
            for idx, generation in enumerate(
                [model_generations[x] for x in ids_require_verification]
            )
        }
        for future in tqdm(as_completed(futures), total=len(ids_require_verification)):
            idx = futures[future]
            model_labels[idx] = future.result()
    return np.array(model_labels)

ft_verif_labels = evaluate(ft_ids_require_verification, ft_model_generations)
ft_verif_labels[:5]

  8%|▊         | 57/747 [00:38<06:08,  1.87it/s]

Cannot parse the answer <think>
1

The original sentence "And what are those risks?" is grammatically correct. All four experts agree that no changes are needed, and the proofreader's correction matches the original. Therefore, the proofreader's correction is valid.. Replacing with 0.

 14%|█▍        | 107/747 [01:20<06:35,  1.62it/s]

Cannot parse the answer <think>
1. Replacing with 0.

 20%|█▉        | 146/747 [01:48<07:02,  1.42it/s]

Cannot parse the answer <think>
1

The original sentence has several errors: "yonger" is misspelled as "younger," the verb "need" should be "needs" for third-person singular, and the phrase "need to more than" is grammatically incorrect. The proofreader's correction addresses all these issues by changing "yonger" to "younger," "need" to "needs," and removing the unnecessary "to." The corrected sentence matches the valid corrections provided by the experts (options 0, 2, 3), making the proofreader's correction valid.. Replacing with 0.

 20%|██        | 150/747 [01:49<03:58,  2.51it/s]

Cannot parse the answer <think>
1. Replacing with 0.

 31%|███       | 230/747 [02:41<05:40,  1.52it/s]

Cannot parse the answer <think>
1. Replacing with 0.

bs_verif_labels = evaluate(bs_ids_require_verification, bs_model_generations)
bs_verif_labels[:5]

 19%|█▉        | 145/747 [01:21<04:47,  2.09it/s]

Cannot parse the answer <think>
1. Replacing with 0.

 25%|██▍       | 184/747 [01:43<03:21,  2.79it/s]

Cannot parse the answer <think>
1

The original sentence "I think it will be lost." is grammatically correct. All four experts agree that no correction is needed. The proofreader's correction matches the original, confirming its validity. Since there are no errors to address, the proofreader's response is correct.. Replacing with 0.

 31%|███▏      | 235/747 [02:08<02:53,  2.96it/s]

Cannot parse the answer <think>
1. Replacing with 0.

 36%|███▌      | 270/747 [02:25<05:30,  1.44it/s]

Cannot parse the answer <think>
1

The original sentence "It can certainly be the case." is grammatically correct. All four experts agreed that no changes were needed. The proofreader's correction matches the original, indicating they found no errors. Since there are no mistakes to correct, the proofreader's response is valid.

1. Replacing with 0.

 59%|█████▉    | 440/747 [03:56<02:02,  2.51it/s]

Cannot parse the answer <think>
1

The original sentence uses "creative" (adjective) and "innovation" (noun), which are mismatched. The correct versions should pair either both nouns ("creativity and innovation") or both adjectives ("creative and innovative"). The proofreader's correction changes "creative" to "creativity" (noun) to match "innovation" (noun), which aligns with the experts' corrections. Therefore, the proofreader's edit is valid.. Replacing with 0.

%%time 

def print_confidence_interval(p, num_obs):
    # Critical value for 95% CI
    z = 1.96
    deviation = z * ((p * (1 - p) / num_obs) ** .5)
    left_boundary = p - deviation
    right_boundary = p + deviation
    print(f'95% confidence interval: {p:.3f} ± {deviation:.3f}\tLeft boundary: {left_boundary:.3f}\tRight boundary: {right_boundary:.3f}')

# Fine-tuned model
ft_verif_dict = {idx: label for idx, label in zip(ft_ids_require_verification, ft_verif_labels)}
ft_results.update(ft_verif_dict)
ft_accuracy = np.mean(list(ft_results.values()))
print('Fine-tuned model:')
print_confidence_interval(ft_accuracy, len(ft_results))

# Baseline model
bs_verif_dict = {idx: label for idx, label in zip(bs_ids_require_verification, bs_verif_labels)}
bs_results.update(bs_verif_dict)
# Coinciding ids
for idx in coinciding_ids:
    bs_results[idx] = ft_verif_dict[idx]
bs_accuracy = np.mean(list(bs_results.values()))
print('Baseline model:')
print_confidence_interval(bs_accuracy, len(bs_results))

assert len(ft_results) == len(bs_results) == len(eval_dataset)

Fine-tuned model:
95% confidence interval: 0.838 ± 0.026	Left boundary: 0.812	Right boundary: 0.864
Baseline model:
95% confidence interval: 0.896 ± 0.022	Left boundary: 0.874	Right boundary: 0.918
CPU times: user 667 μs, sys: 0 ns, total: 667 μs
Wall time: 637 μs

Introduction to Model Distillation: Efficient Knowledge Transfer for AI Applications - Part 3¶

Pre requisites¶

1 - Load dependencies¶

2 - Configuration¶

3 - Initialize Client¶

4 - Test our model¶

5 - Evaluate the model¶

5.1 - Load the dataset.¶

5.2 - Preprocess it by removing spaces before punctuation marks.¶

6 - Create answers¶

7 - Which answers need LLM verification?¶

8- Using LLM for verification¶

9- Measure accuracy of answers¶

Conclusion¶

References¶