Reward Model Training Test: Qwen3-4B¶

# Environment Setup
import os
from dotenv import load_dotenv
load_dotenv()

# Force text-based progress instead of HTML widgets
os.environ["TQDM_NOTEBOOK"] = "false"

# CRITICAL: Import unsloth FIRST for proper TRL patching
import unsloth
from unsloth import is_bf16_supported

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardConfig, RewardTrainer
from peft import LoraConfig, get_peft_model
from datasets import Dataset

# Environment summary
gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
print(f"Environment: unsloth {unsloth.__version__}, PyTorch {torch.__version__}, {gpu}")
print(f"HF_TOKEN loaded: {'Yes' if os.environ.get('HF_TOKEN') else 'No'}")

# Load Qwen3-4B as Reward Model (SequenceClassification head)
from transformers import BitsAndBytesConfig

MODEL_NAME = "Qwen/Qwen3-4B"
print(f"\nLoading {MODEL_NAME.split('/')[-1]} as reward model...")

# 4-bit quantization config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16 if is_bf16_supported() else torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=1,  # Single scalar reward output
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)

# Ensure pad token is set
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id
    model.config.pad_token_id = tokenizer.pad_token_id

print(f"Model loaded: {type(model).__name__}")

# Apply LoRA adapters for reward model training
lora_config = LoraConfig(
    r=16,
    lora_alpha=16,
    lora_dropout=0,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    bias="none",
    task_type="SEQ_CLS",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Create minimal synthetic preference dataset (5 samples)
# Reward models learn from preference pairs (chosen > rejected)

preference_data = [
    {
        "prompt": "Explain recursion in programming.",
        "chosen": "Recursion is when a function calls itself with a simpler version of the problem, including a base case to stop infinite loops.",
        "rejected": "Recursion is just loops."
    },
    {
        "prompt": "What is an API?",
        "chosen": "An API (Application Programming Interface) is a set of protocols that allows different software applications to communicate with each other.",
        "rejected": "API is code."
    },
    {
        "prompt": "Describe version control.",
        "chosen": "Version control is a system that records changes to files over time, allowing you to recall specific versions and collaborate with others.",
        "rejected": "Version control saves files."
    },
    {
        "prompt": "What is a database?",
        "chosen": "A database is an organized collection of structured data stored electronically, typically managed by a database management system (DBMS).",
        "rejected": "A database stores stuff."
    },
    {
        "prompt": "Explain object-oriented programming.",
        "chosen": "Object-oriented programming (OOP) is a paradigm that organizes code into objects containing data (attributes) and behavior (methods).",
        "rejected": "OOP uses objects."
    },
]

dataset = Dataset.from_list(preference_data)
print(f"Dataset created: {len(dataset)} preference pairs")

# Post-training reward scoring test
model.eval()

def get_reward(prompt, response):
    """Score a response using the trained reward model."""
    text = prompt + " " + response
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model(**inputs)
        reward = outputs.logits[0, 0].item()
    
    return reward

# Test with a good and bad response
test_prompt = "What is machine learning?"
good_response = "Machine learning is a subset of AI where computers learn patterns from data to make predictions."
bad_response = "ML is stuff."

good_score = get_reward(test_prompt, good_response)
bad_score = get_reward(test_prompt, bad_response)

print("=" * 60)
print("Reward Model Training Pipeline Test PASSED")
print("=" * 60)
print(f"Good response score: {good_score:.4f}")
print(f"Bad response score:  {bad_score:.4f}")
print(f"Preference correct:  {'Yes' if good_score > bad_score else 'No'}")