[Feature Request] Retrieval Augmented Finetuning Cookbook

[tituslhy]

Hello! I'd like to request a cookbook on retrieval augmented finetuning, please. I've cobbled out some code but I miss the familiarity and the convenience brought about by Unsloth's abstractions. Or can I just use Unsloth's abstractions directly?

Many thanks in advance!

Getting the dataset:

from llama_index.packs.raft_dataset import RAFTDatasetPack
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding

raft_dataset = RAFTDatasetPack(
   file_path = "../data/some_file.pdf",
   llm = Ollama(model="qwen2.5", timeout=120.0),
   embed_model=OllamaEmbedding(model_name="nomic-embed-text")
)

dataset = raft_dataset.run()
output_path = "raft_train"

# Save as .arrow format
dataset.save_to_disk(output_path)

# Save as .jsonl format
dataset.to_json(output_path + ".jsonl")

Beginning the finetuning proper

from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    TrainingArguments
)
from trl import SFTTrainer
from peft import LoraConfig
from multiprocessing import cpu_count
raw_ds = load_dataset("json", data_files="raft_train.jsonl", split="train")


# We only need the two columns:
# - 'instruction': "<DOCUMENT>…</DOCUMENT>…QUESTION…"
# - 'cot_answer': the gold answer text

ds = raw_ds.remove_columns([c for c in raw_ds.column_namesif c not in ["instruction", "cot_answer"]])

# Optionally: split into train/test
splits = ds.train_test_split(test_size=0.1)
train_ds = splits["train"]
eval_ds = splits["test"]

# Initialize tokenizer & model 
MODEL_ID = "meta-llama/Llama-3.2-1B-Instruct"
bnb_cfg = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=True,
   bnb_4bit_compute_dtype="float16"
)
model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     quantization_config=bnb_cfg,
     device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
tokenizer.pad_token = tokenizer.eos_token

def preprocess(batch):
    # batch["instruction"] and batch["cot_answer"] are strings
    inputs = tokenizer(batch["instruction"],truncation=True,max_length=2048,padding=False,)

    # target labels

    with tokenizer.as_target_tokenizer():
          labels = tokenizer(batch["cot_answer"],truncation=True,max_length=512,padding=False,)

     inputs["labels"] = labels["input_ids"]

    return inputs



train_tkn = train_ds.map(preprocess,batched=True,remove_columns=train_ds.column_names,num_proc=cpu_count(),)

eval_tkn = eval_ds.map(preprocess,batched=True,remove_columns=eval_ds.column_names,num_proc=cpu_count(),)

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

training_args = TrainingArguments(
output_dir="raft-sft-output",
per_device_train_batch_size=1, # small batches if quantized
per_device_eval_batch_size=1,
gradient_accumulation_steps=8,
learning_rate=2e-5,
max_steps=60, # or set num_train_epochs
save_strategy="no",
fp16=True,
gradient_checkpointing=True,
logging_steps=5,
seed=42,
optim="adamw_torch",
lr_scheduler_type="cosine",
)

trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=train_tkn,
eval_dataset=eval_tkn,
peft_config=peft_config,
)

trainer.train()

[tituslhy]

I wrote this cookbook using Unsloth! Haven't added in annotations yet but would love some feedback!

Notebook

[rolandtannous]

@tituslhy feel free to share in the discord.

Will convert this to a discussion. Thank you!