!pip install  -U -q transformers trl datasets bitsandbytes peft accelerate
# Tested with transformers==4.46.3, trl==0.12.2, datasets==3.2.0, bitsandbytes==0.45.0, peft==0.14.0, accelerate==1.2.0

!pip install -q flash-attn --no-build-isolation

from huggingface_hub import notebook_login

notebook_login()

from datasets import load_dataset

dataset_id = "HuggingFaceH4/rlaif-v_formatted"
train_dataset, test_dataset = load_dataset(dataset_id, split=['train[:6%]', 'test[:1%]'])

from PIL import Image

def ensure_rgb(example):
    # 이미지가 RGB 형식이 아니라면, RGB로 변환합니다.
    image = example['images'][0]
    if isinstance(image, Image.Image):
        if image.mode != 'RGB':
            image = image.convert('RGB')
        example['images'] = [image]
    return example

# 데이터셋에 변환을 적용합니다.
train_dataset = train_dataset.map(ensure_rgb, num_proc=32)
test_dataset = test_dataset.map(ensure_rgb, num_proc=32)

train_dataset[20]

>>> train_dataset[20]['images'][0]

import torch
from transformers import Idefics3ForConditionalGeneration, AutoProcessor

model_id = "HuggingFaceTB/SmolVLM-Instruct"

from transformers import BitsAndBytesConfig

# BitsAndBytesConfig int-4 설정
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# 모델과 토크나이저 불러오기
model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config,
    _attn_implementation="flash_attention_2",
)
processor = AutoProcessor.from_pretrained(model_id)

>>> from peft import LoraConfig, get_peft_model

>>> # LoRA 구성하기
>>> peft_config = LoraConfig(
...     r=8,
...     lora_alpha=8,
...     lora_dropout=0.1,
...     target_modules=['down_proj','o_proj','k_proj','q_proj','gate_proj','up_proj','v_proj'],
...     use_dora=True,
...     init_lora_weights="gaussian"
... )

>>> # 모델에 PEFT 적용하기
>>> peft_model = get_peft_model(model, peft_config)

>>> # 학습가능한 파라미터 수 출력
>>> peft_model.print_trainable_parameters()

from trl import DPOConfig

training_args = DPOConfig(
    output_dir="smolvlm-instruct-trl-dpo-rlaif-v",
    bf16=True,
    gradient_checkpointing=True,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=32,
    num_train_epochs=5,
    dataset_num_proc=8,  # 토큰화는 8개 프로세스를 사용할 예정입니다.
    dataloader_num_workers=8,  # 데이터를 불러올 때 8개의 worker를 사용할 예정입니다.
    logging_steps=10,
    report_to="tensorboard",
    push_to_hub=True,
    save_strategy="steps",
    save_steps=10,
    save_total_limit=1,
    eval_steps=10,  # 평가를 수행할 스텝 간격
    eval_strategy="steps",
)

from trl import DPOTrainer

trainer = DPOTrainer(
    model=model,
    ref_model=None,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    peft_config=peft_config,
    tokenizer=processor,
)

trainer.train()

trainer.save_model(training_args.output_dir)

>>> import gc
>>> import time

>>> def clear_memory():
...     # 만약 전역 변수들이 있다면 삭제합니다
...     if 'inputs' in globals(): del globals()['inputs']
...     if 'model' in globals(): del globals()['model']
...     if 'processor' in globals(): del globals()['processor']
...     if 'trainer' in globals(): del globals()['trainer']
...     if 'peft_model' in globals(): del globals()['peft_model']
...     if 'bnb_config' in globals(): del globals()['bnb_config']
...     time.sleep(2)

...     # 가비지 컬렉션 수행 및 CUDA 메모리 정리
...     gc.collect()
...     time.sleep(2)
...     torch.cuda.empty_cache()
...     torch.cuda.synchronize()
...     time.sleep(2)
...     gc.collect()
...     time.sleep(2)

...     print(f"GPU allocated memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
...     print(f"GPU reserved memory: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

>>> clear_memory()

model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    _attn_implementation="flash_attention_2",
)

processor = AutoProcessor.from_pretrained(model_id)

adapter_path = "sergiopaniego/smolvlm-instruct-trl-dpo-rlaif-v" # 본인의 Hugging Face 계정 ID를 입력하세요!
model.load_adapter(adapter_path)

test_dataset[20]

>>> test_dataset[20]['images'][0]

def generate_text_from_sample(model, processor, sample, max_new_tokens=1024, device="cuda"):
    # 대화 템플릿을 적용한 텍스트 입력을 준비합니다
    text_input = processor.apply_chat_template(
        sample['prompt'],
        add_generation_prompt=True
    )

    image_inputs = []
    image = sample['images'][0]
    if image.mode != 'RGB':
        image = image.convert('RGB')
    image_inputs.append([image])

    # 모델에 넣을 입력값을 준비합니다
    model_inputs = processor(
        text=text_input,
        images=image_inputs,
        return_tensors="pt",
    ).to(device)  # 특정 device로 입력값을 옮깁니다

    # 모델로 텍스트를 생성합니다
    generated_ids = model.generate(**model_inputs, max_new_tokens=max_new_tokens)

    # 입력값의 인덱스를 제거하기 위해 생성된 인덱스를 잘라줍니다
    trimmed_generated_ids = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    # 출력 텍스트를 디코딩해줍니다
    output_text = processor.batch_decode(
        trimmed_generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False
    )

    return output_text[0]  # 첫 번째로 디코딩된 출력 텍스트를 반환합니다

output = generate_text_from_sample(model, processor, test_dataset[20])
output

from IPython.display import IFrame

IFrame(src="https://sergiopaniego-smolvlm-trl-dpo-rlaif-v.hf.space", width=1000, height=800)