Ajustar o Gemma para tarefas de visão usando o Hugging Face Transformers e o QLoRA

Acessar em ai.google.dev Executar no Google Colab Executar no Kaggle Abrir na Vertex AI Ver código-fonte no GitHub

Este guia mostra como ajustar o Gemma em um conjunto de dados personalizado de imagem e texto para uma tarefa de visão (geração de descrições de produtos) usando os transformadores do Hugging Face Transformers e TRL. Você vai aprender:

  • O que é a adaptação de classificação baixa quantizada (QLoRA)
  • Configurar o ambiente de desenvolvimento
  • Criar e preparar o conjunto de dados de ajuste fino
  • Ajustar o Gemma usando o TRL e o SFTTrainer
  • Testar a inferência do modelo e gerar descrições de produtos com base em imagens e texto.

O que é a adaptação de classificação baixa quantizada (QLoRA)

Este guia demonstra o uso da adaptação de baixa classificação quantizada (QLoRA), que surgiu como um método popular para ajustar LLMs de maneira eficiente, já que reduz os requisitos de recursos computacionais, mantendo o alto desempenho. No QLoRA, o modelo pré-treinado é quantizado para 4 bits e os pesos são congelados. Em seguida, as camadas de adaptador treináveis (LoRA) são anexadas, e apenas as camadas de adaptador são treinadas. Depois disso, os pesos do adaptador podem ser mesclados com o modelo base ou mantidos como um adaptador separado.

Configurar o ambiente de desenvolvimento

A primeira etapa é instalar as bibliotecas do Hugging Face, incluindo o TRL, e os conjuntos de dados para ajustar o modelo aberto.

# Install Pytorch & other libraries
%pip install torch tensorboard torchvision

# Install Transformers
%pip install transformers

# Install Hugging Face libraries
%pip install datasets accelerate evaluate bitsandbytes trl peft protobuf pillow sentencepiece

# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100
#%pip install flash-attn

Observação: se você estiver usando uma GPU com arquitetura Ampere (como a NVIDIA L4) ou mais recente, poderá usar o Flash Attention. O Flash Attention é um método que acelera significativamente os cálculos e reduz o uso da memória de quadrático para linear no comprimento da sequência, o que acelera o treinamento em até 3 vezes. Saiba mais em FlashAttention.

Você precisa de um token válido do Hugging Face para publicar seu modelo. Se você estiver executando em um Google Colab, poderá usar seu token do Hugging Face com segurança usando os segredos do Colab. Caso contrário, defina o token diretamente no método login. Verifique se o token também tem acesso de gravação, já que você envia o modelo para o Hub durante o treinamento.

# Login into Hugging Face Hub
from huggingface_hub import login
login()

Criar e preparar o conjunto de dados de ajuste fino

Ao ajustar LLMs, é importante conhecer seu caso de uso e a tarefa que você quer resolver. Isso ajuda a criar um conjunto de dados para ajustar o modelo. Se você ainda não definiu seu caso de uso, talvez seja necessário voltar ao início.

Como exemplo, este guia se concentra no seguinte caso de uso:

  • Ajustar um modelo Gemma para gerar descrições de produtos concisas e otimizadas para SEO para uma plataforma de e-commerce, especificamente adaptadas para pesquisa em dispositivos móveis.

Este guia usa o conjunto de dados philschmid/amazon-product-descriptions-vlm, um conjunto de dados de descrições de produtos da Amazon, incluindo imagens e categorias de produtos.

O Hugging Face TRL oferece suporte a conversas multimodais. A parte importante é a função "image", que informa à classe de processamento que ela precisa carregar a imagem. A estrutura precisa seguir:

{"messages": [{"role": "system", "content": [{"type": "text", "text":"You are..."}]}, {"role": "user", "content": [{"type": "text", "text": "..."}, {"type": "image"}]}, {"role": "assistant", "content": [{"type": "text", "text": "..."}]}]}
{"messages": [{"role": "system", "content": [{"type": "text", "text":"You are..."}]}, {"role": "user", "content": [{"type": "text", "text": "..."}, {"type": "image"}]}, {"role": "assistant", "content": [{"type": "text", "text": "..."}]}]}
{"messages": [{"role": "system", "content": [{"type": "text", "text":"You are..."}]}, {"role": "user", "content": [{"type": "text", "text": "..."}, {"type": "image"}]}, {"role": "assistant", "content": [{"type": "text", "text": "..."}]}]}

Agora você pode usar a biblioteca de conjuntos de dados do Hugging Face para carregar o conjunto de dados e criar um modelo de comando para combinar a imagem, o nome do produto e a categoria e adicionar uma mensagem do sistema. O conjunto de dados inclui imagens como objetosPil.Image.

from datasets import load_dataset
from PIL import Image

# System message for the assistant
system_message = "You are an expert product description writer for Amazon."

# User prompt that combines the user query and the schema
user_prompt = """Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.
Only return description. The description should be SEO optimized and for a better mobile search experience.

<PRODUCT>
{product}
</PRODUCT>

<CATEGORY>
{category}
</CATEGORY>
"""

# Convert dataset to OAI messages
def format_data(sample):
    return {
        "messages": [
            {
                "role": "system",
                #"content": [{"type": "text", "text": system_message}],
                "content": system_message,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_prompt.format(
                            product=sample["Product Name"],
                            category=sample["Category"],
                        ),
                    },
                    {
                        "type": "image",
                        "image": sample["image"],
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": sample["description"]}],
            },
        ],
    }

def process_vision_info(messages: list[dict]) -> list[Image.Image]:
    image_inputs = []
    # Iterate through each conversation
    for msg in messages:
        # Get content (ensure it's a list)
        content = msg.get("content", [])
        if not isinstance(content, list):
            content = [content]

        # Check each content element for images
        for element in content:
            if isinstance(element, dict) and (
                "image" in element or element.get("type") == "image"
            ):
                # Get the image and convert to RGB
                if "image" in element:
                    image = element["image"]
                else:
                    image = element
                image_inputs.append(image.convert("RGB"))
    return image_inputs

# Load dataset from the hub
dataset = load_dataset("philschmid/amazon-product-descriptions-vlm", split="train")
dataset = dataset.train_test_split(test_size=0.1)

# Convert dataset to OAI messages
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
dataset_train = [format_data(sample) for sample in dataset["train"]]
dataset_test = [format_data(sample) for sample in dataset["test"]]

print(dataset_train[345]["messages"])

README.md: 0.00B [00:00, ?B/s]
data/train-00000-of-00001.parquet:   0%|          | 0.00/47.6M [00:00<?, ?B/s]
Generating train split:   0%|          | 0/1345 [00:00<?, ? examples/s]
[{'role': 'system', 'content': 'You are an expert product description writer for Amazon.'}, {'role': 'user', 'content': [{'type': 'text', 'text': "Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.\nOnly return description. The description should be SEO optimized and for a better mobile search experience.\n\n<PRODUCT>\nRazor Agitator BMX/Freestyle Bike, 20-Inch\n</PRODUCT>\n\n<CATEGORY>\nSports & Outdoors | Outdoor Recreation | Cycling | Kids' Bikes & Accessories | Kids' Bikes\n</CATEGORY>\n"}, {'type': 'image', 'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=500x413 at 0x7B7250181790>}]}, {'role': 'assistant', 'content': [{'type': 'text', 'text': 'Conquer the streets with the Razor Agitator BMX Bike! This 20-inch freestyle bike is built for young riders ready to take on any challenge. Durable frame, responsive handling – perfect for tricks and cruising.  Get yours today!'}]}]

Ajustar o Gemma usando o TRL e o SFTTrainer

Agora está tudo pronto para ajustar o modelo. O Hugging Face TRL SFTTrainer facilita a supervisão do ajuste fino de LLMs abertos. O SFTTrainer é uma subclasse do Trainer da biblioteca transformers e oferece suporte a todos os mesmos recursos, incluindo registro, avaliação e checkpoint, mas adiciona outros recursos de qualidade de vida, incluindo:

  • Formatação de conjuntos de dados, incluindo formatos conversacionais e de instruções
  • Treinamento apenas em conclusões, ignorando comandos
  • Empacotamento de conjuntos de dados para um treinamento mais eficiente
  • Suporte a ajustes finos com eficiência de parâmetros (PEFT), incluindo QLoRA
  • Preparação do modelo e do tokenizador para ajuste fino conversacional (como a adição de tokens especiais)

O código a seguir carrega o modelo Gemma e o tokenizador do Hugging Face e inicializa a configuração de quantização.

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig

# Hugging Face model id
model_id = "google/gemma-4-E2B" # @param ["google/gemma-4-E2B","google/gemma-4-E4B"] {"allow-input":true}

# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] < 8:
    raise ValueError("GPU does not support bfloat16, please use a GPU that supports bfloat16.")

# Define model init arguments
model_kwargs = dict(
    dtype=torch.bfloat16, # What torch dtype to use, defaults to auto
    device_map="auto", # Let torch decide how to load the model
)

# BitsAndBytesConfig int-4 config
model_kwargs["quantization_config"] = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=model_kwargs["dtype"],
    bnb_4bit_quant_storage=model_kwargs["dtype"],
)

# Load model and tokenizer
model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
processor = AutoProcessor.from_pretrained("google/gemma-4-E2B-it") # Load the Instruction Tokenizer to use the official Gemma template

config.json: 0.00B [00:00, ?B/s]
model.safetensors:   0%|          | 0.00/10.2G [00:00<?, ?B/s]
Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
generation_config.json:   0%|          | 0.00/149 [00:00<?, ?B/s]
processor_config.json: 0.00B [00:00, ?B/s]
chat_template.jinja: 0.00B [00:00, ?B/s]
config.json: 0.00B [00:00, ?B/s]
tokenizer_config.json: 0.00B [00:00, ?B/s]
tokenizer.json:   0%|          | 0.00/32.2M [00:00<?, ?B/s]

O SFTTrainer oferece suporte a uma integração integrada com peft, o que facilita o ajuste eficiente de LLMs usando QLoRA. Basta criar um LoraConfig e fornecê-lo ao treinador.

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    modules_to_save=["lm_head", "embed_tokens"], # make sure to save the lm_head and embed_tokens as you train the special tokens
    ensure_weight_tying=True,
)

Antes de iniciar o treinamento, é necessário definir o hiperparâmetro que você quer usar em um SFTConfig e um collate_fn personalizado para processar a visão. O collate_fn converte as mensagens com texto e imagens em um formato que o modelo possa entender.

from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-product-description",     # directory to save and repository id
    num_train_epochs=3,                         # number of training epochs
    per_device_train_batch_size=1,              # batch size per device during training
    optim="adamw_torch_fused",                  # use fused adamw optimizer
    logging_steps=5,                            # log every 5 steps
    save_strategy="epoch",                      # save checkpoint every epoch
    eval_strategy="epoch",                      # evaluate checkpoint every epoch
    learning_rate=2e-4,                         # learning rate, based on QLoRA paper
    bf16=True,                                  # use bfloat16 precision
    max_grad_norm=0.3,                          # max gradient norm based on QLoRA paper
    lr_scheduler_type="constant",               # use constant learning rate scheduler
    push_to_hub=True,                           # push model to hub
    report_to="tensorboard",                    # report metrics to tensorboard
    dataset_text_field="",                      # need a dummy field for collator
    dataset_kwargs={"skip_prepare_dataset": True}, # important for collator
    remove_unused_columns = False               # important for collator
)

# Create a data collator to encode text and image pairs
def collate_fn(examples):
    texts = []
    images = []
    for example in examples:
        image_inputs = process_vision_info(example["messages"])
        text = processor.apply_chat_template(
            example["messages"], add_generation_prompt=False, tokenize=False
        )
        texts.append(text.strip())
        images.append(image_inputs)

    # Tokenize the texts and process the images
    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)

    # The labels are the input_ids, and we mask the padding tokens and image tokens in the loss computation
    labels = batch["input_ids"].clone()

    # Mask tokens for not being used in the loss computation
    labels[labels == processor.tokenizer.pad_token_id] = -100
    labels[labels == processor.tokenizer.boi_token_id] = -100
    labels[labels == processor.tokenizer.image_token_id] = -100
    labels[labels == processor.tokenizer.eoi_token_id] = -100

    batch["labels"] = labels
    return batch

Agora você tem todos os elementos básicos necessários para criar seu SFTTrainer e iniciar o treinamento do modelo.

from trl import SFTTrainer

# Create Trainer object
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset_train,
    eval_dataset=dataset_test,
    peft_config=peft_config,
    processing_class=processor,
    data_collator=collate_fn,
)

Para iniciar o treinamento, chame o método train().

# Start training, the model will be automatically saved to the Hub and the output directory
trainer.train()

# Save the final model again to the Hugging Face Hub
trainer.save_model()

The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'eos_token_id': 1, 'bos_token_id': 2, 'pad_token_id': 0}.

Antes de testar o modelo, libere a memória.

# free the memory again
del model
del trainer
torch.cuda.empty_cache()

Ao usar o QLoRA, você só treina adaptadores e não o modelo completo. Isso significa que, ao salvar o modelo durante o treinamento, você só salva os pesos do adaptador e não o modelo completo. Se você quiser salvar o modelo completo, o que facilita o uso com pilhas de veiculação como vLLM ou TGI, poderá mesclar os pesos do adaptador nos pesos do modelo usando o método merge_and_unload e, em seguida, salvar o modelo com o método save_pretrained. Isso salva um modelo padrão, que pode ser usado para inferência.

from peft import PeftModel

# Load Model base model
model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)

# Merge LoRA and base model and save
peft_model = PeftModel.from_pretrained(model, args.output_dir)
merged_model = peft_model.merge_and_unload()
merged_model.save_pretrained("merged_model", safe_serialization=True, max_shard_size="2GB")

processor = AutoProcessor.from_pretrained(args.output_dir)
processor.save_pretrained("merged_model")

Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
Writing model shards:   0%|          | 0/5 [00:00<?, ?it/s]
['merged_model/processor_config.json']

Testar a inferência do modelo e gerar descrições de produtos

Depois que o treinamento for concluído, avalie e teste o modelo. É possível carregar diferentes amostras do conjunto de dados de teste e avaliar o modelo nessas amostras.

model_id = "merged_model"

# Load Model with PEFT adapter
model = AutoModelForImageTextToText.from_pretrained(
  model_id,
  device_map="auto",
  dtype="auto",
)
processor = AutoProcessor.from_pretrained(model_id)

Loading weights:   0%|          | 0/2012 [00:00<?, ?it/s]
The tied weights mapping and config for this model specifies to tie model.language_model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints with different values, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning.

Você pode testar a inferência fornecendo um nome, categoria e imagem do produto. A sample inclui uma figura de ação da Marvel.

import requests
from PIL import Image

# Test sample with Product Name, Category and Image
sample = {
  "product_name": "Hasbro Marvel Avengers-Serie Marvel Assemble Titan-Held, Iron Man, 30,5 cm Actionfigur",
  "category": "Toys & Games | Toy Figures & Playsets | Action Figures",
  "image": Image.open(requests.get("https://m.media-amazon.com/images/I/81+7Up7IWyL._AC_SY300_SX300_.jpg", stream=True).raw).convert("RGB")
}

def generate_description(sample, model, processor):
    # Convert sample into messages and then apply the chat template
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": [
            {"type": "image","image": sample["image"]},
            {"type": "text", "text": user_prompt.format(product=sample["product_name"], category=sample["category"])},
        ]},
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    print(text)
    # Process the image and text
    image_inputs = process_vision_info(messages)
    # Tokenize the text and process the images
    inputs = processor(
        text=[text],
        images=image_inputs,
        padding=True,
        return_tensors="pt",
    )
    # Move the inputs to the device
    inputs = inputs.to(model.device)

    # Generate the output
    stop_token_ids = [processor.tokenizer.eos_token_id, processor.tokenizer.convert_tokens_to_ids("<turn|>")]
    generated_ids = model.generate(**inputs, max_new_tokens=256, top_p=1.0, do_sample=True, temperature=0.8, eos_token_id=stop_token_ids, disable_compile=True)
    # Trim the generation and decode the output to text
    generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return output_text[0]

# generate the description
description = generate_description(sample, model, processor)
print("MODEL OUTPUT>> \n")
print(description)

<bos><|turn>system
You are an expert product description writer for Amazon.<turn|>
<|turn>user


<|image|>

Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.
Only return description. The description should be SEO optimized and for a better mobile search experience.

<PRODUCT>
Hasbro Marvel Avengers-Serie Marvel Assemble Titan-Held, Iron Man, 30,5 cm Actionfigur
</PRODUCT>

<CATEGORY>
Toys & Games | Toy Figures & Playsets | Action Figures
</CATEGORY><turn|>
<|turn>model

MODEL OUTPUT>> 

Enhance your collection with the Marvel Avengers - Avengers Assemble Ultron-Comforter Set! This soft and cuddly blanket and pillowcase feature everyone's favorite Avengers, Iron Man, and his loyal companion War Machine. Officially licensed by Marvel.  Bring home the heroic team!

Resumo e próximas etapas

Este tutorial abordou como ajustar um modelo Gemma para tarefas de visão usando o TRL e o QLoRA, especificamente para gerar descrições de produtos. Confira os seguintes documentos: