使用 Hugging Face Transformers 與 QLoRA 微調 Gemma¶

# Install Pytorch & other libraries
%pip install torch tensorboard
%pip install -U torchao

# Install Transformers
%pip install "transformers>=5.10.1"

# Install Hugging Face libraries
%pip install datasets accelerate evaluate bitsandbytes trl "peft>=0.19.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

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

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, idx):
  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.select(range(1250))

# Convert dataset to OAI messages
dataset = dataset.map(create_conversation, with_indices=True, remove_columns=dataset.features)
# split dataset into 80% training samples and 20% test samples
dataset = dataset.train_test_split(test_size=0.2, shuffle=False)

# Print formatted user prompt
for item in dataset["train"][0]["messages"]:
  print(item)

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    # no target_modules --- PEFT's Gemma 4 defaults scope to the LM layers
    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,
)

import torch
from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-text-to-sql",         # directory to save and repository id
    max_length=512,                         # max length for model and packing of the dataset
    num_train_epochs=3,                     # number of training epochs
    per_device_train_batch_size=1,          # batch size per device during training
    per_device_eval_batch_size=1,           # batch size per device during evaluation
    optim="adamw_torch_fused",              # use fused adamw optimizer
    logging_steps=10,                       # log every 10 steps
    save_strategy="epoch",                  # save checkpoint every epoch
    eval_strategy="epoch",                  # evaluate checkpoint every epoch
    learning_rate=2e-4,                     # learning rate
    fp16=True if torch_dtype == torch.float16 else False,  # use float16 precision
    bf16=True if torch_dtype == torch.bfloat16 else False, # use bfloat16 precision
    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={"skip_prepare_dataset": True}, # important for collator
    remove_unused_columns = False,                 # important for collator
)

# Data collator
def collate_fn(examples):
    texts = []

    for example in examples:
        full_text = processor.apply_chat_template(
            example["messages"], add_generation_prompt=False, tokenize=False
        )
        texts.append(full_text.strip())

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

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

    target_tokens = [
        processor.tokenizer.convert_tokens_to_ids("<|turn>"),
        processor.tokenizer.convert_tokens_to_ids("model"),
        processor.tokenizer.convert_tokens_to_ids("\n")
    ]
    target_len = len(target_tokens)

    for i in range(labels.size(0)):
        row_tokens = batch["input_ids"][i].tolist()

        # Find where the assistant block begins
        assistant_start_idx = None
        for idx in range(len(row_tokens) - target_len + 1):
            if row_tokens[idx : idx + target_len] == target_tokens:
                # We want to keep loss calculation on the assistant transcription tokens,
                # so we move the index right past the assistant header ('<|turn>\nmodel\n')
                assistant_start_idx = idx + target_len
                break

        if assistant_start_idx is not None:
            # Mask everything from index 0 up to the start of the actual Japanese text response
            labels[i, :assistant_start_idx] = -100
        else:
            # Fallback safety: if template matching fails for an anomalous row, mask padding anyway
            print("WARNING: maybe the sample is too long, try to increase `token_limit` value.")
            labels[i, labels[i] == processor.tokenizer.pad_token_id] = -100

        """
        # --- DEBUG PRINT CODE ---
        print(f"\n--- Example {i} (Split index: {assistant_start_idx}) ---")
        debug_string = []
        for token_id, label_id in zip(row_tokens, labels[i].tolist()):
            # Decode token by token so we can see exactly what is masked
            decoded_token = processor.tokenizer.decode([token_id])

            if label_id == -100:
                # Red text for masked tokens (ANSI Escape Code)
                debug_string.append(f"\033[91m{decoded_token}\033[0m")
            else:
                # Green text for active loss tokens
                debug_string.append(f"\033[92m{decoded_token}\033[0m")

        print("".join(debug_string))
        # ---
        """


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

    batch["labels"] = labels
    return batch

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,
)

# 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()

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

from transformers import AutoModelForMultimodalLM, AutoProcessor

model_id = "merged_model"

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

from random import randint
import re
from transformers import pipeline, GenerationConfig, pipeline

config = GenerationConfig.from_pretrained(model_id)
config.max_new_tokens = 256
config.eos_token_id = [processor.tokenizer.convert_tokens_to_ids("<turn|>")]

# Load the model and tokenizer into the pipeline
pipe = pipeline("text-generation", model=model, tokenizer=processor.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
prompt = processor.tokenizer.apply_chat_template(test_sample["messages"][:2], tokenize=False, add_generation_prompt=True)
print(prompt)

# Generate our SQL query.
outputs = pipe(text_inputs=prompt, generation_config=config)

# Extract the user query and original answer
print(f"Context:\n", re.search(r'<SCHEMA>\n(.*?)\n</SCHEMA>', test_sample['messages'][1]['content'], re.DOTALL).group(1).strip())
print(f"Query:\n", re.search(r'<USER_QUERY>\n(.*?)\n</USER_QUERY>', test_sample['messages'][1]['content'], re.DOTALL).group(1).strip())
print(f"Original Answer:\n{test_sample['messages'][2]['content']}")
print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip().removesuffix("<turn|>")}")