Fine-Tuning TinyLLaMA (1.1B) untuk Mengubah Deskripsi Teks Menjadi Warna

Fine-tuning adalah proses melatih kembali model AI yang sudah dilatih sebelumnya (pretrained model) dengan data baru yang lebih spesifik agar model lebih sesuai dengan kebutuhan atau domain tertentu.

Fine-tuning ini menggunakan Google Colab dengan dataset burkelibbey/colors dari hugging face.

Ekspektasi fine-tuning ini adalah model bisa memahami dan memberikan jawaban warna dalam bentuk Hex Code berdasarkan deskripsi teks.

Contohnya: Prompt: “Give me a pure brown color” Response: #804421

Tools:

1. Base model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
2. Dataset: burkelibbey/colors
3. LoRA (PEFT) untuk fine-tuning yang efisien
4. Quantization 4-bit (BitsAndBytes) agar bisa jalan di Colab
5. Library: transformers, trl, peft, datasets

Install Library:

!pip install accelerate peft bitsandbytes transformers trl

Full Code Fine-Tuning:

import torch
from datasets import load_dataset, Dataset
from peft import LoraConfig, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
from trl import SFTTrainer
import os

dataset="burkelibbey/colors"
base_model="TinyLlama/TinyLlama-1.1B-Chat-v1.0"
output_model="tinyllama-colorist-v1"

def prepare_train_data(data_id):
    data = load_dataset(data_id, split="train")
    data_df = data.to_pandas()
    data_df["text"] = data_df[["description", "color"]].apply(lambda x: "<|user|>\n" + x["description"] + "</s>\n<|assistant|>\n" + x["color"] + "</s>", axis=1)
    data = Dataset.from_pandas(data_df)
    return data

def get_model_and_tokenizer(mode_id):

    tokenizer = AutoTokenizer.from_pretrained(mode_id)
    tokenizer.pad_token = tokenizer.eos_token
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype="float16", bnb_4bit_use_double_quant=True
    )
    model = AutoModelForCausalLM.from_pretrained(
        mode_id, quantization_config=bnb_config, device_map="auto"
    )
    model.config.use_cache=False
    model.config.pretraining_tp=1
    return model, tokenizer

def finetune_tinyllama(dataset, base_model, output_model):
    data = prepare_train_data(dataset)
    peft_config = LoraConfig(
            r=8, lora_alpha=16, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
        )

    model, tokenizer = get_model_and_tokenizer(base_model)

    training_arguments = TrainingArguments(
            output_dir=output_model,
            per_device_train_batch_size=16,
            gradient_accumulation_steps=4,
            optim="paged_adamw_32bit",
            learning_rate=2e-4,
            lr_scheduler_type="cosine",
            save_strategy="epoch",
            logging_steps=10,
            num_train_epochs=3,
            max_steps=250,
            fp16=True,
            report_to="none"
            # push_to_hub=True
        )

    trainer = SFTTrainer(
        model=model,
        train_dataset=data,
        peft_config=peft_config,
        args=training_arguments,
        processing_class=tokenizer,
    )

    trainer.train()

if __name__ == "__main__":
    finetune_tinyllama(dataset, base_model, output_model)

Penjelasan:

Import Libraries

import torch
from datasets import load_dataset, Dataset
from peft import LoraConfig, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
from trl import SFTTrainer
import os

Mengimpor semua library penting:

• torch: library utama untuk komputasi tensor dan model.
• datasets: dari Hugging Face, untuk load dataset seperti "burkelibbey/colors".
• peft: untuk konfigurasi LoRA (Parameter Efficient Fine-Tuning).
• transformers: untuk load model/tokenizer dari Hugging Face.
• trl: library dari Hugging Face untuk SFT (Supervised Fine-Tuning).
• os: manipulasi file system jika dibutuhkan.

Inisialisasi Parameter

dataset="burkelibbey/colors"
base_model="TinyLlama/TinyLlama-1.1B-Chat-v1.0"
output_model="tinyllama-colorist-v1"

• dataset: ID dataset dari Hugging Face Hub.
• base_model: model TinyLLaMA dari Hugging Face yang akan difine-tune.
• output_model: nama direktori untuk menyimpan model hasil fine-tuning.

def prepare_train_data(data_id):
    data = load_dataset(data_id, split="train")
    data_df = data.to_pandas()
    data_df["text"] = data_df[["description", "color"]].apply(lambda x: "<|user|>\n" + x["description"] + "</s>\n<|assistant|>\n" + x["color"] + "</s>", axis=1)
    data = Dataset.from_pandas(data_df)
    return data

Fungsi: prepare_train_data(data_id)

Tujuan: Menyiapkan data pelatihan dalam format teks dialog antara user dan AI (Chat-like prompt-response) untuk fine-tuning model generatif seperti TinyLLaMA.

1. Load dataset dari Hugging Face dengan ID data_id.
2. Ubah ke DataFrame agar mudah dimanipulasi.
3. Gabungkan description dan color menjadi satu kolom text dalam format prompt-response:

• <|user|> description </s> <|assistant|> color </s>

4. Konversi kembali ke Hugging Face Dataset.

5. Return data siap latih.

def get_model_and_tokenizer(mode_id):

    tokenizer = AutoTokenizer.from_pretrained(mode_id)
    tokenizer.pad_token = tokenizer.eos_token
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype="float16", bnb_4bit_use_double_quant=True
    )
    model = AutoModelForCausalLM.from_pretrained(
        mode_id, quantization_config=bnb_config, device_map="auto"
    )
    model.config.use_cache=False
    model.config.pretraining_tp=1
    return model, tokenizer

Fungsi get_model_and_tokenizer memuat model dan tokenizer dari Hugging Face (mode_id) dengan konfigurasi quantization 4-bit agar hemat memori (cocok untuk Colab). Tokenizer disiapkan dan model diatur agar siap untuk fine-tuning

def finetune_tinyllama(dataset, base_model, output_model):
    data = prepare_train_data(dataset)
    peft_config = LoraConfig(
            r=8, lora_alpha=16, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM"
        )

    model, tokenizer = get_model_and_tokenizer(base_model)

    training_arguments = TrainingArguments(
            output_dir=output_model,
            per_device_train_batch_size=16,
            gradient_accumulation_steps=4,
            optim="paged_adamw_32bit",
            learning_rate=2e-4,
            lr_scheduler_type="cosine",
            save_strategy="epoch",
            logging_steps=10,
            num_train_epochs=3,
            max_steps=250,
            fp16=True,
            report_to="none"
            # push_to_hub=True
        )

    trainer = SFTTrainer(
        model=model,
        train_dataset=data,
        peft_config=peft_config,
        args=training_arguments,
        processing_class=tokenizer,
    )

    trainer.train()

ungsi finetune_tinyllama melatih model TinyLLaMA dengan teknik LoRA menggunakan data percakapan. Intinya:

• Siapkan data.
• Atur konfigurasi LoRA.
• Muat model dan tokenizer.
• Tentukan parameter training.
• Jalankan training dengan SFTTrainer.

Konfigurasi training:

batch_size=16 → Model memproses 16 data sekaligus per langkah. Dipilih kecil agar cocok dengan memori GPU di Google Colab.

gradient_accumulation_steps=4 → Gradien dari 4 batch (masing-masing 16 data) dikumpulkan dulu sebelum update bobot. Ini setara batch size 64, tanpa butuh memori besar.

fp16=True → hemat memori (pakai float16).

logging_steps=10 → Logging setiap 10 step.

max_steps=250 → Total hanya 250 langkah training dilakukan. Cocok untuk eksperimen cepat, melatih model tanpa memakan waktu lama atau kehabisan resource di Colab. Note: Hapus max_steps jika ingin fine-tuning seluruh data

num_train_epochs=3 → Epoch adalah 1 kali proses model melihat seluruh data pelatihan dari awal sampai akhir.

Jadi kalau:

• Dataset berisi 1.000 data
• Dan batch size = 100 → Maka 1 epoch = 10 langkah (steps)

if __name__ == "__main__":
    finetune_tinyllama(dataset, base_model, output_model)

Execute fine-tuning finetune_tinyllama

Merging Lora Dengan Base Model:

Menggabungkan hasil fine-tune (LoRA adapter) ke model utama dan menyimpan model akhir dalam bentuk standalone (tidak tergantung adapter lagi).

from peft import AutoPeftModelForCausalLM, PeftModel
from transformers import AutoModelForCausalLM
import torch
import os

save_dir = "tinyllama-colorist-v1-merged"
model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.float16, load_in_8bit=False,
                                             device_map="auto",
                                             trust_remote_code=True)

model_path = "/content/tinyllama-colorist-v1/checkpoint-250"

peft_model = PeftModel.from_pretrained(model, model_path, from_transformers=True, device_map="auto")

model = peft_model.merge_and_unload()

# Save merged model
model.save_pretrained(save_dir)

# Save tokenizer from base model to same folder
tokenizer = AutoTokenizer.from_pretrained(base_model)
tokenizer.pad_token = tokenizer.eos_token  
tokenizer.save_pretrained(save_dir)

Penjelasan:

save_dir = "tinyllama-colorist-v1-merged"

• Lokasi direktori tempat model final (yang sudah digabung) akan disimpan.

model_path = "/content/tinyllama-colorist-v1/checkpoint-250"

• Lokasi hasil fine-tuning (adapter LoRA) pada langkah ke-250.

peft_model = PeftModel.from_pretrained(model, model_path, from_transformers=True, device_map="auto")

• Memuat adapter LoRA (peft_model) dan menggabungkannya dengan model dasar.

model = peft_model.merge_and_unload()

• Menggabungkan bobot dari adapter LoRA ke model utama (merge), lalu membuang adapter (unload).
• Hasilnya: model sudah standalone, siap disimpan dan digunakan tanpa file LoRA lagi.

model.save_pretrained(save_dir)

• Menyimpan model hasil merge ke folder tinyllama-colorist-v1-merged.

tokenizer = AutoTokenizer.from_pretrained(base_model)
tokenizer.pad_token = tokenizer.eos_token  
tokenizer.save_pretrained(save_dir)

• Memuat tokenizer dari model dasar dan menyimpan ke folder yang sama, agar bisa digunakan bersama model final.

Hasil Akhir:

Folder tinyllama-colorist-v1-merged berisi:

• Model hasil fine-tune yang sudah digabung
• Tokenizer
• Siap digunakan untuk inferensi di mana saja, tanpa perlu file adapter lagi.

Inference Model

Kode ini digunakan untuk memuat model hasil fine-tune dan mengujinya dengan memberikan prompt, lalu mencetak respons yang dihasilkan oleh model.

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from time import perf_counter

# Load dari folder hasil merge
model_path = "tinyllama-colorist-v1-merged"

tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto"
)
model.eval()

def format_prompt(user_input):
    return f"<|user|>\n{user_input}</s>\n<|assistant|>"

def generate_response(user_input):
    prompt = format_prompt(user_input)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    generation_config = {
        "do_sample": True,
        "top_k": 5,
        "temperature": 0.5,
        "repetition_penalty": 1.2,
        "max_new_tokens": 12,
        "pad_token_id": tokenizer.eos_token_id
    }

    start = perf_counter()
    with torch.no_grad():
        outputs = model.generate(**inputs, **generation_config)
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = decoded.replace(prompt, "").strip()

    print("Assistant:", response)
    print(f"⏱ Time: {round(perf_counter() - start, 2)}s")

# Contoh penggunaan
generate_response("give me a pure brown color")

Penjelasan Key Code

model_path = "tinyllama-colorist-v1-merged"

• Path ke folder model yang sudah digabung (LoRA + base model).

tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token

• Memuat tokenizer dari folder model.
• Menyetel token padding sama dengan EOS token agar tidak error saat batching.

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto"
)
model.eval()

• Memuat model dari folder, dengan dtype float16 untuk efisiensi.
• device_map="auto" akan otomatis memilih GPU jika tersedia.
• model.eval() menonaktifkan training mode (misalnya dropout), penting untuk inferensi.

def format_prompt(user_input):
    return f"<|user|>\n{user_input}</s>\n<|assistant|>"

• Fungsi untuk membentuk prompt yang sesuai dengan format fine-tuning:
• <|user|> give me a pure brown color</s> <|assistant|>

def generate_response(user_input):
    prompt = format_prompt(user_input)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

• Menerima input dari user → ubah jadi prompt → tokenisasi → pindah ke device (GPU/CPU).

generation_config = {
        "do_sample": True,             # Sampling aktif (bukan greedy decoding)
        "top_k": 5,                    # Ambil dari 5 token teratas untuk variasi
        "temperature": 0.5,            # Kontrol kreativitas (lebih rendah = lebih deterministik)
        "repetition_penalty": 1.2,     # Hindari pengulangan output
        "max_new_tokens": 12,          # Output maksimal 12 token
        "pad_token_id": tokenizer.eos_token_id
    }

• Pengaturan untuk proses text generation.

decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = decoded.replace(prompt, "").strip()

• Decode output ke teks biasa.
• Hilangkan prompt dari hasil, sehingga hanya sisa jawaban model.

print("Assistant:", response)
    print(f"⏱ Time: {round(perf_counter() - start, 2)}s")

• Cetak hasil dan waktu inferensi.

Contoh Output:

generate_response("give me a pure brown color")

Output:

🟢 Assistant: #804421
⏱ Time: 0.56s

Lihat Hasil Akhir di Google Colab: https://colab.research.google.com/drive/1MoRzTYyBfTA9X_X3mm9IsI8MVjMJhi14#scrollTo=2mGIHSY_jos8