Instruction Tuning in Machine Learning

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig

# Load base model and tokenizer
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
)

# Load instruction dataset (Alpaca format)
dataset = load_dataset("tatsu-lab/alpaca", split="train")

# Format into chat template
def format_instruction(example):
    if example["input"]:
        text = f"""### Instruction:\n{example['instruction']}\n\n### Input:\n{example['input']}\n\n### Response:\n{example['output']}"""
    else:
        text = f"""### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"""
    return {"text": text}

dataset = dataset.map(format_instruction)

# LoRA configuration for memory efficiency
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    task_type="CAUSAL_LM",
)

# Training configuration
training_args = SFTConfig(
    output_dir="./instruction-tuned-llama",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,
    logging_steps=10,
    save_strategy="epoch",
    max_seq_length=2048,
    dataset_text_field="text",
    packing=True,  # Pack short examples for efficiency
)

# Train
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    peft_config=lora_config,
    tokenizer=tokenizer,
)
trainer.train()
trainer.save_model("./instruction-tuned-llama-final")

import openai
import json
import random
from typing import List, Dict

client = openai.OpenAI()  # Uses OPENAI_API_KEY env var

# Seed tasks (8-10 human-written examples to bootstrap)
SEED_TASKS = [
    {"instruction": "Write a haiku about machine learning.", "input": "", "output": "Neurons learn and grow\nGradients flowing like streams\nPatterns emerge clear"},
    {"instruction": "Classify the sentiment of the following review.", "input": "The food was absolutely terrible and the service was slow.", "output": "Negative"},
    {"instruction": "Explain the concept of gradient descent in simple terms.", "input": "", "output": "Gradient descent is like finding the lowest point in a valley while blindfolded. You feel the slope under your feet and take a step downhill. You repeat this process, and eventually you reach the bottom."},
]

def generate_instructions(seed_tasks: List[Dict], n_generate: int = 20) -> List[Dict]:
    """Generate new instruction-response pairs using self-instruct."""
    generated = []
    task_pool = seed_tasks.copy()
    
    for i in range(n_generate):
        # Sample 3 examples from the pool as few-shot demonstrations
        demos = random.sample(task_pool, min(3, len(task_pool)))
        demo_text = "\n\n".join([
            f"Instruction: {d['instruction']}\n"
            f"Input: {d['input']}\n"
            f"Output: {d['output']}"
            for d in demos
        ])
        
        prompt = f"""Below are example instruction-input-output triples. Generate a NEW, diverse instruction that is different from the examples. The instruction should be clear, specific, and answerable.

{demo_text}

Now generate a new example in JSON format with keys: instruction, input, output.
Make the instruction different in topic and format from the examples above."""
        
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.9,
            max_tokens=1024,
        )
        
        try:
            new_task = json.loads(response.choices[0].message.content)
            if all(k in new_task for k in ["instruction", "input", "output"]):
                generated.append(new_task)
                task_pool.append(new_task)  # Add to pool for diversity
                print(f"Generated {i+1}/{n_generate}: {new_task['instruction'][:60]}...")
        except (json.JSONDecodeError, KeyError):
            continue  # Skip malformed outputs
    
    return generated

# Generate 1000 synthetic instruction pairs
synthetic_data = generate_instructions(SEED_TASKS, n_generate=1000)

# Save for training
with open("synthetic_instructions.json", "w") as f:
    json.dump(synthetic_data, f, indent=2)

print(f"Generated {len(synthetic_data)} instruction pairs")
print(f"Estimated API cost: ~${len(synthetic_data) * 0.005:.2f}")

from transformers import AutoTokenizer
from typing import List, Dict
import torch

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

def format_sharegpt_with_loss_mask(
    conversation: List[Dict[str, str]],
    tokenizer,
    max_length: int = 4096,
) -> Dict[str, torch.Tensor]:
    """
    Format a ShareGPT-style multi-turn conversation with proper loss masking.
    Only compute loss on assistant responses, not user messages or system prompt.
    """
    # Build the full conversation string with role markers
    input_ids = []
    labels = []
    
    for turn in conversation:
        role = turn["from"]  # "system", "human", or "gpt"
        content = turn["value"]
        
        if role == "human":
            text = f"[INST] {content} [/INST]"
            tokens = tokenizer.encode(text, add_special_tokens=False)
            input_ids.extend(tokens)
            labels.extend([-100] * len(tokens))  # Mask user tokens
            
        elif role == "gpt":
            text = f"{content}</s>"
            tokens = tokenizer.encode(text, add_special_tokens=False)
            input_ids.extend(tokens)
            labels.extend(tokens)  # Compute loss on assistant tokens
            
        elif role == "system":
            text = f"<<SYS>>\n{content}\n<</SYS>>\n\n"
            tokens = tokenizer.encode(text, add_special_tokens=False)
            input_ids.extend(tokens)
            labels.extend([-100] * len(tokens))  # Mask system tokens
    
    # Truncate to max length
    input_ids = input_ids[:max_length]
    labels = labels[:max_length]
    
    return {
        "input_ids": torch.tensor(input_ids),
        "labels": torch.tensor(labels),
        "attention_mask": torch.ones(len(input_ids), dtype=torch.long),
    }

# Example ShareGPT conversation
conversation = [
    {"from": "system", "value": "You are a helpful AI assistant."},
    {"from": "human", "value": "What is instruction tuning?"},
    {"from": "gpt", "value": "Instruction tuning is a fine-tuning technique where a pretrained language model is trained on instruction-response pairs to learn to follow natural language instructions."},
    {"from": "human", "value": "How does it differ from RLHF?"},
    {"from": "gpt", "value": "Instruction tuning uses supervised learning on demonstrations, while RLHF uses reinforcement learning with a reward model trained on human preferences. Instruction tuning typically comes first in the training pipeline."},
]

result = format_sharegpt_with_loss_mask(conversation, tokenizer)
print(f"Sequence length: {len(result['input_ids'])}")
print(f"Tokens with loss: {(result['labels'] != -100).sum().item()}")
print(f"Masked tokens: {(result['labels'] == -100).sum().item()}")

import openai
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Tuple

client = openai.OpenAI()

JUDGE_PROMPT = """Rate the quality of the following instruction-response pair on a scale of 1-5.

Criteria:
- Instruction clarity (is the instruction specific and unambiguous?)
- Response accuracy (is the response correct and complete?)
- Response helpfulness (does it actually address the instruction?)
- Response quality (is it well-written, well-structured?)

Instruction: {instruction}
Response: {response}

Respond with ONLY a JSON object: {{"score": <1-5>, "reason": "<brief explanation>"}}"""

def score_example(example: Dict) -> Tuple[Dict, int, str]:
    """Score a single instruction-response pair using GPT-4o-mini as judge."""
    prompt = JUDGE_PROMPT.format(
        instruction=example["instruction"],
        response=example["output"],
    )
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
            max_tokens=200,
        )
        result = json.loads(response.choices[0].message.content)
        return example, result["score"], result["reason"]
    except Exception:
        return example, 0, "scoring_failed"

def filter_dataset(
    dataset: List[Dict],
    min_score: int = 4,
    max_workers: int = 10,
) -> List[Dict]:
    """Filter instruction dataset, keeping only high-quality examples."""
    filtered = []
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(score_example, ex): ex for ex in dataset}
        
        for future in as_completed(futures):
            example, score, reason = future.result()
            if score >= min_score:
                example["quality_score"] = score
                filtered.append(example)
    
    print(f"Kept {len(filtered)}/{len(dataset)} examples (score >= {min_score})")
    print(f"Filter rate: {1 - len(filtered)/len(dataset):.1%}")
    return filtered

# Usage
raw_dataset = json.load(open("synthetic_instructions.json"))
filtered_dataset = filter_dataset(raw_dataset, min_score=4)
json.dump(filtered_dataset, open("filtered_instructions.json", "w"), indent=2)

# Axolotl configuration for instruction tuning (YAML)
base_model: meta-llama/Llama-2-7b-hf
model_type: LlamaForCausalLM
tokenizer_type: LlamaTokenizer

load_in_8bit: false
load_in_4bit: true       # QLoRA

adapter: lora
lora_r: 16
lora_alpha: 32
lora_dropout: 0.05
lora_target_modules:
  - q_proj
  - k_proj
  - v_proj
  - o_proj

datasets:
  - path: tatsu-lab/alpaca
    type: alpaca
  - path: Open-Orca/OpenOrca
    type: sharegpt
    conversation: chatml

dataset_prepared_path: ./prepared_data
val_set_size: 0.02

sequence_len: 4096
sample_packing: true
pad_to_sequence_len: true

wandb_project: instruction-tuning
wandb_run_id: llama2-7b-sft-v1

gradient_accumulation_steps: 4
micro_batch_size: 2
num_epochs: 3
learning_rate: 2e-4
lr_scheduler: cosine
warmup_ratio: 0.03
optimizer: adamw_torch
weight_decay: 0.01

bf16: true
tf32: true
gradient_checkpointing: true

save_strategy: epoch
evaluation_strategy: steps
eval_steps: 100