scripts/create_notebook.py - main - dany/cu-scholarship-red

scripts/create_notebook.py
652 строки20 KB
dany
Add task 1 SFT style training notebook section
15 июл 2026, 09:38
15 июл 2026, 09:381d5114c
#!/usr/bin/env python3
import json
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "notebooks" / "red_alignment_pipeline.ipynb"


def md(source):
    return {"cell_type": "markdown", "metadata": {}, "source": source.strip() + "\n"}


def code(source):
    return {
        "cell_type": "code",
        "execution_count": None,
        "metadata": {},
        "outputs": [],
        "source": source.strip() + "\n",
    }


SETUP_CELLS = [
    md(
        """
        # Red Level ML Olympiad: Alignment Pipeline

        This Colab notebook trains all adapters inside the notebook and prints the five public metrics.

        Fixed rules:
        - `seed = 42`
        - greedy generation for style metrics
        - length-normalized response log-probability for implicit-preference metrics
        """
    ),
    code(
        """
        !pip -q install -r requirements.txt
        """
    ),
    code(
        """
        import gc
        import json
        import math
        import os
        import pickle
        import random
        from pathlib import Path

        import numpy as np
        import torch
        from datasets import Dataset
        from scipy.sparse import hstack
        from sklearn.metrics import accuracy_score
        from transformers import (
            AutoModelForCausalLM,
            AutoModelForSequenceClassification,
            AutoTokenizer,
            BitsAndBytesConfig,
            set_seed,
        )
        from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training

        try:
            from trl import (
                DPOConfig,
                DPOTrainer,
                RewardConfig,
                RewardTrainer,
                SFTConfig,
                SFTTrainer,
            )
        except ImportError as exc:
            raise RuntimeError("Install requirements.txt first.") from exc

        SEED = 42
        random.seed(SEED)
        np.random.seed(SEED)
        torch.manual_seed(SEED)
        torch.cuda.manual_seed_all(SEED)
        set_seed(SEED)

        ROOT = Path.cwd()
        DATA_DIR = ROOT / "data"
        METRICS_DIR = ROOT / "metrics"
        ARTIFACTS = ROOT / "artifacts"
        ARTIFACTS.mkdir(exist_ok=True)

        MODEL_NAME = "Qwen/Qwen3-4B-Instruct-2507"
        MAX_SEQ_LENGTH = 768
        MAX_NEW_TOKENS = 180
        """
    ),
    code(
        """
        def read_jsonl(path):
            with open(path, encoding="utf-8") as f:
                return [json.loads(line) for line in f]


        kid_adult = read_jsonl(DATA_DIR / "kid_adult.jsonl")
        good_bad = read_jsonl(DATA_DIR / "good_bad.jsonl")
        public_test_style = read_jsonl(DATA_DIR / "public_test_style.jsonl")
        public_test_quality = read_jsonl(DATA_DIR / "public_test_quality.jsonl")

        print(len(kid_adult), len(good_bad), len(public_test_style), len(public_test_quality))
        """
    ),
    code(
        """
        with open(METRICS_DIR / "style_clf.pkl", "rb") as f:
            style_obj = pickle.load(f)

        style_clf = style_obj["clf"]
        style_vecs = style_obj["vecs"]
        simple_class_index = list(style_clf.classes_).index(1)


        def p_simple(texts):
            X = hstack([vec.transform(texts) for vec in style_vecs])
            return style_clf.predict_proba(X)[:, simple_class_index]


        print("Reference public kid P_simple:", float(np.mean(p_simple([x["kid"] for x in public_test_style]))))
        print("Reference public adult P_simple:", float(np.mean(p_simple([x["adult"] for x in public_test_style]))))
        """
    ),
    code(
        """
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )

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


        def load_tokenizer():
            tok = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
            if tok.pad_token is None:
                tok.pad_token = tok.eos_token
            tok.padding_side = "right"
            return tok


        tokenizer = load_tokenizer()


        def user_prompt(prompt):
            return tokenizer.apply_chat_template(
                [{"role": "user", "content": prompt}],
                tokenize=False,
                add_generation_prompt=True,
            )


        def supervised_text(prompt, answer):
            return tokenizer.apply_chat_template(
                [
                    {"role": "user", "content": prompt},
                    {"role": "assistant", "content": answer},
                ],
                tokenize=False,
                add_generation_prompt=False,
            )


        def free_memory(*objs):
            for obj in objs:
                del obj
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        """
    ),
    code(
        """
        @torch.no_grad()
        def generate_answers(model, records, prompt_key="prompt", max_new_tokens=MAX_NEW_TOKENS):
            model.eval()
            answers = []
            for row in records:
                prompt = user_prompt(row[prompt_key])
                inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
                out = model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    do_sample=False,
                    temperature=None,
                    top_p=None,
                    pad_token_id=tokenizer.pad_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                )
                gen = out[0, inputs["input_ids"].shape[1] :]
                answers.append(tokenizer.decode(gen, skip_special_tokens=True).strip())
            return answers


        def print_interval(name, value, bins):
            for label, lo, hi in bins:
                if (lo is None or value >= lo) and (hi is None or value < hi or (hi == 1.0 and value <= hi)):
                    print(f"{name}: {value:.4f} -> {label}")
                    return label
            raise ValueError(value)


        STYLE_BINS = [
            ("< 0.4", None, 0.4),
            ("0.4 - 0.7", 0.4, 0.7),
            ("0.7 - 0.9", 0.7, 0.9),
            ("0.9 - 1.0", 0.9, 1.0),
        ]
        QUALITY_BINS = [
            ("< 0.6", None, 0.6),
            ("0.6 - 0.75", 0.6, 0.75),
            ("0.75 - 0.9", 0.75, 0.9),
            ("0.9 - 1.0", 0.9, 1.0),
        ]
        """
    ),
]


TASK1_CELLS = [
    md("## Task 1 - SFT style transfer"),
    code(
        """
        def load_causal_model():
            model = AutoModelForCausalLM.from_pretrained(
                MODEL_NAME,
                quantization_config=bnb_config,
                device_map="auto",
                trust_remote_code=True,
            )
            model.config.use_cache = False
            model = prepare_model_for_kbit_training(model)
            return model


        sft_train = Dataset.from_list(
            [{"text": supervised_text(row["prompt"], row["kid"])} for row in kid_adult]
        )

        sft_model = load_causal_model()
        sft_args = SFTConfig(
            output_dir=str(ARTIFACTS / "task1_sft"),
            num_train_epochs=1,
            per_device_train_batch_size=1,
            gradient_accumulation_steps=8,
            learning_rate=2e-4,
            logging_steps=20,
            save_strategy="epoch",
            seed=SEED,
            bf16=True,
            max_seq_length=MAX_SEQ_LENGTH,
            dataset_text_field="text",
            report_to=[],
        )
        sft_trainer = SFTTrainer(
            model=sft_model,
            args=sft_args,
            train_dataset=sft_train,
            tokenizer=tokenizer,
            peft_config=lora_config,
        )
        sft_trainer.train()
        sft_trainer.model.save_pretrained(ARTIFACTS / "task1_sft_adapter")
        tokenizer.save_pretrained(ARTIFACTS / "task1_sft_adapter")
        """
    ),
    code(
        """
        task1_answers = generate_answers(sft_trainer.model, public_test_style)
        task1_p_simple = float(np.mean(p_simple(task1_answers)))
        task1_interval = print_interval("Task 1 P_simple", task1_p_simple, STYLE_BINS)
        """
    ),
]


TASK2_CELLS = [
    md("## Task 2 - DPO style preference"),
    code(
        """
        dpo_style_train = Dataset.from_list(
            [
                {
                    "prompt": user_prompt(row["prompt"]),
                    "chosen": row["kid"],
                    "rejected": row["adult"],
                }
                for row in kid_adult
            ]
        )

        dpo_style_args = DPOConfig(
            output_dir=str(ARTIFACTS / "task2_dpo_style"),
            num_train_epochs=1,
            per_device_train_batch_size=1,
            gradient_accumulation_steps=8,
            learning_rate=8e-5,
            beta=0.1,
            logging_steps=20,
            save_strategy="epoch",
            seed=SEED,
            bf16=True,
            max_length=MAX_SEQ_LENGTH,
            max_prompt_length=384,
            report_to=[],
        )
        dpo_style_trainer = DPOTrainer(
            model=sft_trainer.model,
            ref_model=None,
            args=dpo_style_args,
            train_dataset=dpo_style_train,
            tokenizer=tokenizer,
        )
        dpo_style_trainer.train()
        dpo_style_trainer.model.save_pretrained(ARTIFACTS / "task2_dpo_style_adapter")
        tokenizer.save_pretrained(ARTIFACTS / "task2_dpo_style_adapter")
        """
    ),
    code(
        """
        task2_answers = generate_answers(dpo_style_trainer.model, public_test_style)
        task2_p_simple = float(np.mean(p_simple(task2_answers)))
        task2_interval = print_interval("Task 2 P_simple", task2_p_simple, STYLE_BINS)
        """
    ),
]


TASK3_CELLS = [
    md("## Task 3 - Reward model for quality"),
    code(
        """
        reward_tokenizer = tokenizer
        rm_lora_config = LoraConfig(
            r=16,
            lora_alpha=32,
            lora_dropout=0.05,
            bias="none",
            task_type="SEQ_CLS",
            target_modules=[
                "q_proj",
                "k_proj",
                "v_proj",
                "o_proj",
                "gate_proj",
                "up_proj",
                "down_proj",
            ],
        )

        reward_model = AutoModelForSequenceClassification.from_pretrained(
            MODEL_NAME,
            num_labels=1,
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
        )
        reward_model.config.pad_token_id = reward_tokenizer.pad_token_id
        reward_model.config.use_cache = False
        reward_model = prepare_model_for_kbit_training(reward_model)

        rm_train = Dataset.from_list(
            [
                {
                    "chosen": supervised_text(row["instruction"], row["chosen"]),
                    "rejected": supervised_text(row["instruction"], row["rejected"]),
                }
                for row in good_bad
            ]
        )

        rm_args = RewardConfig(
            output_dir=str(ARTIFACTS / "task3_reward_model"),
            num_train_epochs=1,
            per_device_train_batch_size=1,
            gradient_accumulation_steps=8,
            learning_rate=1e-4,
            logging_steps=20,
            save_strategy="epoch",
            seed=SEED,
            bf16=True,
            max_length=MAX_SEQ_LENGTH,
            report_to=[],
        )
        rm_trainer = RewardTrainer(
            model=reward_model,
            args=rm_args,
            tokenizer=reward_tokenizer,
            train_dataset=rm_train,
            peft_config=rm_lora_config,
        )
        rm_trainer.train()
        rm_trainer.model.save_pretrained(ARTIFACTS / "task3_reward_adapter")
        reward_tokenizer.save_pretrained(ARTIFACTS / "task3_reward_adapter")
        """
    ),
    code(
        """
        @torch.no_grad()
        def reward_score(texts, batch_size=8):
            rm_trainer.model.eval()
            scores = []
            for i in range(0, len(texts), batch_size):
                batch = texts[i : i + batch_size]
                inputs = reward_tokenizer(
                    batch,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=MAX_SEQ_LENGTH,
                ).to(rm_trainer.model.device)
                logits = rm_trainer.model(**inputs).logits.squeeze(-1)
                scores.extend(logits.detach().float().cpu().tolist())
            return np.array(scores)


        rm_good = [supervised_text(row["prompt"], row["chosen"]) for row in public_test_quality]
        rm_bad = [supervised_text(row["prompt"], row["rejected"]) for row in public_test_quality]
        task3_accuracy = float(np.mean(reward_score(rm_good) > reward_score(rm_bad)))
        task3_interval = print_interval("Task 3 reward pairwise accuracy", task3_accuracy, QUALITY_BINS)
        """
    ),
]


TASK4_CELLS = [
    md("## Task 4 - DPO quality preference"),
    code(
        """
        dpo_quality_train = Dataset.from_list(
            [
                {
                    "prompt": user_prompt(row["instruction"]),
                    "chosen": row["chosen"],
                    "rejected": row["rejected"],
                }
                for row in good_bad
            ]
        )

        quality_model = dpo_style_trainer.model
        dpo_quality_args = DPOConfig(
            output_dir=str(ARTIFACTS / "task4_dpo_quality"),
            num_train_epochs=1,
            per_device_train_batch_size=1,
            gradient_accumulation_steps=8,
            learning_rate=5e-5,
            beta=0.1,
            logging_steps=20,
            save_strategy="epoch",
            seed=SEED,
            bf16=True,
            max_length=MAX_SEQ_LENGTH,
            max_prompt_length=384,
            report_to=[],
        )
        dpo_quality_trainer = DPOTrainer(
            model=quality_model,
            ref_model=None,
            args=dpo_quality_args,
            train_dataset=dpo_quality_train,
            tokenizer=tokenizer,
        )
        dpo_quality_trainer.train()
        dpo_quality_trainer.model.save_pretrained(ARTIFACTS / "task4_dpo_quality_adapter")
        tokenizer.save_pretrained(ARTIFACTS / "task4_dpo_quality_adapter")
        """
    ),
    code(
        """
        @torch.no_grad()
        def mean_answer_logprob(model, prompt, answer):
            model.eval()
            prompt_text = user_prompt(prompt)
            answer_text = answer + tokenizer.eos_token
            prompt_ids = tokenizer(prompt_text, add_special_tokens=False).input_ids
            full_ids = tokenizer(prompt_text + answer_text, add_special_tokens=False).input_ids
            input_ids = torch.tensor([full_ids], device=model.device)
            labels = input_ids.clone()
            labels[:, : len(prompt_ids)] = -100
            outputs = model(input_ids=input_ids)
            logits = outputs.logits[:, :-1, :]
            shifted_labels = labels[:, 1:]
            mask = shifted_labels.ne(-100)
            token_logprobs = torch.log_softmax(logits, dim=-1)
            selected = token_logprobs.gather(-1, shifted_labels.clamp_min(0).unsqueeze(-1)).squeeze(-1)
            return float((selected[mask].sum() / mask.sum()).detach().cpu())


        def implicit_preference_accuracy(model, rows):
            good_scores = []
            bad_scores = []
            for row in rows:
                good_scores.append(mean_answer_logprob(model, row["prompt"], row["chosen"]))
                bad_scores.append(mean_answer_logprob(model, row["prompt"], row["rejected"]))
            good_scores = np.array(good_scores)
            bad_scores = np.array(bad_scores)
            return float(np.mean(good_scores > bad_scores))


        task4_accuracy = implicit_preference_accuracy(dpo_quality_trainer.model, public_test_quality)
        task4_interval = print_interval("Task 4 DPO implicit-preference accuracy", task4_accuracy, QUALITY_BINS)
        """
    ),
]


TASK5_CELLS = [
    md("## Task 5 - SimPO quality preference"),
    code(
        """
        simpo_model = load_causal_model()
        simpo_sft_adapter = ARTIFACTS / "task2_dpo_style_adapter"
        simpo_model = PeftModel.from_pretrained(simpo_model, simpo_sft_adapter, is_trainable=True)

        try:
            from trl import SimPOConfig, SimPOTrainer
            simpo_args = SimPOConfig(
                output_dir=str(ARTIFACTS / "task5_simpo_quality"),
                num_train_epochs=1,
                per_device_train_batch_size=1,
                gradient_accumulation_steps=8,
                learning_rate=5e-5,
                beta=2.0,
                gamma_beta_ratio=0.5,
                logging_steps=20,
                save_strategy="epoch",
                seed=SEED,
                bf16=True,
                max_length=MAX_SEQ_LENGTH,
                max_prompt_length=384,
                report_to=[],
            )
            simpo_trainer = SimPOTrainer(
                model=simpo_model,
                args=simpo_args,
                train_dataset=dpo_quality_train,
                tokenizer=tokenizer,
            )
        except ImportError:
            simpo_args = DPOConfig(
                output_dir=str(ARTIFACTS / "task5_simpo_quality"),
                num_train_epochs=1,
                per_device_train_batch_size=1,
                gradient_accumulation_steps=8,
                learning_rate=5e-5,
                beta=2.0,
                loss_type="simpo",
                logging_steps=20,
                save_strategy="epoch",
                seed=SEED,
                bf16=True,
                max_length=MAX_SEQ_LENGTH,
                max_prompt_length=384,
                report_to=[],
            )
            simpo_trainer = DPOTrainer(
                model=simpo_model,
                ref_model=None,
                args=simpo_args,
                train_dataset=dpo_quality_train,
                tokenizer=tokenizer,
            )

        simpo_trainer.train()
        simpo_trainer.model.save_pretrained(ARTIFACTS / "task5_simpo_quality_adapter")
        tokenizer.save_pretrained(ARTIFACTS / "task5_simpo_quality_adapter")
        """
    ),
    code(
        """
        task5_accuracy = implicit_preference_accuracy(simpo_trainer.model, public_test_quality)
        task5_interval = print_interval("Task 5 SimPO implicit-preference accuracy", task5_accuracy, QUALITY_BINS)

        print("\\nSummary")
        print({
            "task1_p_simple": task1_p_simple,
            "task1_interval": task1_interval,
            "task2_p_simple": task2_p_simple,
            "task2_interval": task2_interval,
            "task3_accuracy": task3_accuracy,
            "task3_interval": task3_interval,
            "task4_accuracy": task4_accuracy,
            "task4_interval": task4_interval,
            "task5_accuracy": task5_accuracy,
            "task5_interval": task5_interval,
        })
        """
    ),
]


def main():
    stage = int(sys.argv[1]) if len(sys.argv) > 1 else 5
    cells = list(SETUP_CELLS)
    if stage >= 1:
        cells += TASK1_CELLS
    if stage >= 2:
        cells += TASK2_CELLS
    if stage >= 3:
        cells += TASK3_CELLS
    if stage >= 4:
        cells += TASK4_CELLS
    if stage >= 5:
        cells += TASK5_CELLS

    notebook = {
        "cells": cells,
        "metadata": {
            "accelerator": "GPU",
            "colab": {"provenance": []},
            "kernelspec": {
                "display_name": "Python 3",
                "language": "python",
                "name": "python3",
            },
            "language_info": {"name": "python", "version": "3.10"},
        },
        "nbformat": 4,
        "nbformat_minor": 5,
    }
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(notebook, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"Wrote {OUT} with stage {stage}")


if __name__ == "__main__":
    main()