Este guia ensina como ajustar o Gemma em um conjunto de dados personalizado de texto para SQL usando os Transformers e o TRL do Hugging Face. Você vai aprender a realizar as tarefas a seguir:
- O que é a adaptação quantizada de baixa classificação (QLoRA)
- Configurar o ambiente de desenvolvimento
- Criar e preparar o conjunto de dados de ajuste de detalhes
- Ajustar o Gemma usando o TRL e o SFTTrainer
- Testar a inferência de modelo e gerar consultas SQL
O que é a adaptação quantizada de baixa classificação (QLoRA)
Este guia demonstra o uso da adaptação quantizada de baixa classificação (QLoRA), que surgiu como um método conhecido para ajustar LLMs de maneira eficiente, já que reduz os requisitos de recursos computacionais e mantém o alto desempenho. Na QloRA, o modelo pré-treinado é quantizado para 4 bits, e os pesos são congelados. Em seguida, as camadas de adaptador treináveis (LoRA, na sigla em inglês) são anexadas, e apenas as camadas de adaptador são treinadas. Depois, 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, incluindo diferentes técnicas de RLHF e alinhamento.
# Install Pytorch & other libraries
%pip install "torch>=2.4.0" tensorboard
# Install Gemma release branch from Hugging Face
%pip install "transformers>=4.51.3"
# Install Hugging Face libraries
%pip install --upgrade \
"datasets==3.3.2" \
"accelerate==1.4.0" \
"evaluate==0.4.3" \
"bitsandbytes==0.45.3" \
"trl==0.15.2" \
"peft==0.14.0" \
protobuf \
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. A atenção instantânea é um método que acelera significativamente as computações e reduz o uso de memória de quadrático para linear no comprimento da sequência, acelerando o treinamento em até três vezes. Saiba mais em FlashAttention.
Antes de começar o treinamento, é necessário aceitar os Termos de Uso do Gemma. Você pode aceitar a licença no Hugging Face clicando no botão "Concordo e acesse o repositório" na página do modelo em: http://huggingface.co/google/gemma-3-1b-pt
Depois de aceitar a licença, você vai precisar de um token do Hugging Face válido para acessar o 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.
from google.colab import userdata
from huggingface_hub import login
# Login into Hugging Face Hub
hf_token = userdata.get('HF_TOKEN') # If you are running inside a Google Colab
login(hf_token)
Criar e preparar o conjunto de dados de ajuste de detalhes
Ao ajustar os LLMs, é importante conhecer seu caso de uso e a tarefa que você quer resolver. Isso ajuda a criar um conjunto de dados para ajustar seu modelo. Se você ainda não definiu seu caso de uso, talvez seja melhor voltar para a prancheta.
Como exemplo, este guia se concentra no seguinte caso de uso:
- Aprimore uma linguagem natural para o modelo SQL e integre-a perfeitamente a uma ferramenta de análise de dados. O objetivo é reduzir significativamente o tempo e a experiência necessários para gerar consultas SQL, permitindo que até mesmo usuários não técnicos extraiam insights significativos dos dados.
A conversão de texto em SQL pode ser um bom caso de uso para ajustar LLMs, já que é uma tarefa complexa que exige muito conhecimento (interno) sobre os dados e a linguagem SQL.
Depois de determinar que o ajuste fino é a solução certa, você precisa de um conjunto de dados para fazer isso. O conjunto de dados precisa ser diversificado e demonstrar as tarefas que você quer resolver. Há várias maneiras de criar esse conjunto de dados, incluindo:
- Usando conjuntos de dados de código aberto, como o Spider
- Usar conjuntos de dados sintéticos criados por LLMs, como o Alpaca
- Usar conjuntos de dados criados por humanos, como Dolly.
- Usar uma combinação de métodos, como o Orca
Cada um dos métodos tem vantagens e desvantagens e depende do orçamento, do tempo e dos requisitos de qualidade. Por exemplo, usar um conjunto de dados já existente é mais fácil, mas pode não ser adequado para seu caso de uso específico. Já o uso de especialistas de domínio pode ser mais preciso, mas é demorado e caro. Também é possível combinar vários métodos para criar um conjunto de dados de instruções, conforme mostrado em Orca: Progressive Learning from Complex Explanation Traces of GPT-4 (em inglês).
Este guia usa um conjunto de dados já existente (philschmid/gretel-synthetic-text-to-sql), um conjunto de dados sintéticos de Texto para SQL de alta qualidade, incluindo instruções de linguagem natural, definições de esquema, raciocínio e a consulta SQL correspondente.
O TRL do Hug Face oferece suporte à criação automática de modelos de formatos de conjuntos de dados de conversa. Isso significa que você só precisa converter seu conjunto de dados nos objetos JSON corretos, e o trl cuida da criação de modelos e da conversão para o formato certo.
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
O philschmid/gretel-synthetic-text-to-sql (em inglês) contém mais de 100 mil amostras. Para manter o guia pequeno, ele é reduzido para usar apenas 10.000 amostras.
Agora você pode usar a biblioteca Hugging Face Datasets para carregar o conjunto de dados e criar um modelo de comando para combinar a instrução de linguagem natural, a definição do esquema e adicionar uma mensagem do sistema ao seu assistente.
from datasets import load_dataset
# System message for the assistant
system_message = """You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."""
# User prompt that combines the user query and the schema
user_prompt = """Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.
<SCHEMA>
{context}
</SCHEMA>
<USER_QUERY>
{question}
</USER_QUERY>
"""
def create_conversation(sample):
return {
"messages": [
# {"role": "system", "content": system_message},
{"role": "user", "content": user_prompt.format(question=sample["sql_prompt"], context=sample["sql_context"])},
{"role": "assistant", "content": sample["sql"]}
]
}
# Load dataset from the hub
dataset = load_dataset("philschmid/gretel-synthetic-text-to-sql", split="train")
dataset = dataset.shuffle().select(range(12500))
# Convert dataset to OAI messages
dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
# split dataset into 10,000 training samples and 2,500 test samples
dataset = dataset.train_test_split(test_size=2500/12500)
# Print formatted user prompt
print(dataset["train"][345]["messages"][1]["content"])
Ajustar o Gemma usando o TRL e o SFTTrainer
Agora está tudo pronto para ajustar seu modelo. O SFTTrainer do Hugging Face TRL facilita a supervisão do ajuste fino de LLMs abertos. O SFTTrainer é uma subclasse do Trainer da biblioteca transformers e oferece os mesmos recursos, incluindo registro, avaliação e checkpoint, mas adiciona recursos de qualidade de vida, incluindo:
- Formatação de conjuntos de dados, incluindo formatos de conversa e instruções
- Treinamento apenas para conclusões, ignorando instruções
- Empacotar conjuntos de dados para um treinamento mais eficiente
- Suporte a ajustes finos com eficiência de parâmetros (PEFT, na sigla em inglês), incluindo QloRA
- Preparar o modelo e o tokenizer para o ajuste fino de conversação (como adicionar tokens especiais)
O código abaixo carrega o modelo e o tokenizer do Gemma do Hugging Face e inicializa a configuração de quantização.
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, BitsAndBytesConfig
# Hugging Face model id
model_id = "google/gemma-3-1b-pt" # or `google/gemma-3-4b-pt`, `google/gemma-3-12b-pt`, `google/gemma-3-27b-pt`
# Select model class based on id
if model_id == "google/gemma-3-1b-pt":
model_class = AutoModelForCausalLM
else:
model_class = AutoModelForImageTextToText
# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] >= 8:
torch_dtype = torch.bfloat16
else:
torch_dtype = torch.float16
# Define model init arguments
model_kwargs = dict(
attn_implementation="eager", # Use "flash_attention_2" when running on Ampere or newer GPU
torch_dtype=torch_dtype, # What torch dtype to use, defaults to auto
device_map="auto", # Let torch decide how to load the model
)
# BitsAndBytesConfig: Enables 4-bit quantization to reduce model size/memory usage
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['torch_dtype'],
bnb_4bit_quant_storage=model_kwargs['torch_dtype'],
)
# Load model and tokenizer
model = model_class.from_pretrained(model_id, **model_kwargs)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it") # Load the Instruction Tokenizer to use the official Gemma template
O SFTTrainer oferece suporte a uma integração nativa com peft, o que facilita o ajuste eficiente de LLMs usando a QLoRA. Você só precisa criar uma LoraConfig e fornecer ao instrutor.
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
)
Antes de iniciar o treinamento, você precisa definir o hiperparâmetro que quer usar em uma instância SFTConfig.
from trl import SFTConfig
args = SFTConfig(
output_dir="gemma-text-to-sql", # directory to save and repository id
max_seq_length=512, # max sequence length for model and packing of the dataset
packing=True, # Groups multiple samples in the dataset into a single sequence
num_train_epochs=3, # number of training epochs
per_device_train_batch_size=1, # batch size per device during training
gradient_accumulation_steps=4, # number of steps before performing a backward/update pass
gradient_checkpointing=True, # use gradient checkpointing to save memory
optim="adamw_torch_fused", # use fused adamw optimizer
logging_steps=10, # log every 10 steps
save_strategy="epoch", # save checkpoint every epoch
learning_rate=2e-4, # learning rate, based on QLoRA paper
fp16=True if torch_dtype == torch.float16 else False, # use float16 precision
bf16=True if torch_dtype == torch.bfloat16 else False, # use bfloat16 precision
max_grad_norm=0.3, # max gradient norm based on QLoRA paper
warmup_ratio=0.03, # warmup ratio 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_kwargs={
"add_special_tokens": False, # We template with special tokens
"append_concat_token": True, # Add EOS token as separator token between examples
}
)
Agora você tem todos os elementos básicos necessários para criar o SFTTrainer e iniciar o treinamento do modelo.
from trl import SFTTrainer
# Create Trainer object
trainer = SFTTrainer(
model=model,
args=args,
train_dataset=dataset["train"],
peft_config=peft_config,
processing_class=tokenizer
)
Inicie o treinamento chamando 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()
Antes de testar o modelo, libere a memória.
# free the memory again
del model
del trainer
torch.cuda.empty_cache()
Ao usar a QLoRA, você treina apenas os 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, mescle os pesos do adaptador nos pesos do modelo usando o método merge_and_unload e salve 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 = model_class.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 = AutoTokenizer.from_pretrained(args.output_dir)
processor.save_pretrained("merged_model")
Testar a inferência de modelo e gerar consultas SQL
Depois do treinamento, é necessário avaliar e testar o modelo. É possível carregar amostras diferentes do conjunto de dados de teste e avaliar o modelo nessas amostras.
import torch
from transformers import pipeline
model_id = "gemma-text-to-sql"
# Load Model with PEFT adapter
model = model_class.from_pretrained(
model_id,
device_map="auto",
torch_dtype=torch_dtype,
attn_implementation="eager",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
Vamos carregar uma amostra aleatória do conjunto de dados de teste e gerar um comando SQL.
from random import randint
import re
# Load the model and tokenizer into the pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
# Load a random sample from the test dataset
rand_idx = randint(0, len(dataset["test"]))
test_sample = dataset["test"][rand_idx]
# Convert as test example into a prompt with the Gemma template
stop_token_ids = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<end_of_turn>")]
prompt = pipe.tokenizer.apply_chat_template(test_sample["messages"][:2], tokenize=False, add_generation_prompt=True)
# Generate our SQL query.
outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=stop_token_ids, disable_compile=True)
# Extract the user query and original answer
print(f"Context:\n", re.search(r'<SCHEMA>\n(.*?)\n</SCHEMA>', test_sample['messages'][0]['content'], re.DOTALL).group(1).strip())
print(f"Query:\n", re.search(r'<USER_QUERY>\n(.*?)\n</USER_QUERY>', test_sample['messages'][0]['content'], re.DOTALL).group(1).strip())
print(f"Original Answer:\n{test_sample['messages'][1]['content']}")
print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")
Resumo e próximas etapas
Neste tutorial, abordamos como ajustar um modelo Gemma usando TRL e QLoRA. Confira os seguintes documentos:
- Saiba como gerar texto com um modelo Gemma.
- Saiba como ajustar o Gemma para tarefas de visão usando o Hugging Face Transformers.
- Saiba como realizar ajuste fino e inferência distribuídos em um modelo do Gemma.
- Saiba como usar modelos abertos do Gemma com a Vertex AI.
- Aprenda a ajustar o Gemma usando o KerasNLP e implantar na Vertex AI.