Fine-Tuning SLMs with LoRA for Reliable JSON Generation

pip install transformers torch accelerate peft bitsandbytes datasets

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "microsoft/phi-3-mini-4k-instruct"

# Configuration for 4-bit quantization
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# Load the model with quantization
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    device_map="auto", # Automatically map model layers to available devices
    trust_remote_code=True, 
)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token # Set pad token to end-of-sentence token
tokenizer.padding_side = 'right' # Pad on the right to avoid issues with generation

print(f"Model loaded on: {model.device}")

{
  "username": "string",
  "email": "string or null",
  "category": "enum('login_issue', 'billing_problem', 'feature_request', 'bug_report')",
  "priority": "integer(1-5)",
  "summary": "string"
}

import json
import random
from datasets import Dataset

# Predefined templates and options
categories = ['login_issue', 'billing_problem', 'feature_request', 'bug_report']
summaries = [
    "I can't log into my account, it says 'invalid credentials'. My username is {user}.",
    "My recent invoice seems incorrect. Can you check the charges for user {user}? My email is {email}.",
    "I'd love to have an API for exporting my data. This would be a game-changer.",
    "The dashboard crashes when I click the 'Analytics' tab. This is urgent, my username is {user}."
]
users = [("jdoe", "[email protected]"), ("test_user", None), ("alice_b", "[email protected]")]

def generate_example():
    user, email = random.choice(users)
    category = random.choice(categories)
    priority = random.randint(1, 5)
    summary_template = random.choice(summaries)
    
    # Inject user details into the summary text
    input_text = summary_template.format(user=user, email=email or "not provided")
    
    # Create the target JSON output
    target_json = {
        "username": user,
        "email": email,
        "category": category,
        "priority": priority,
        "summary": input_text
    }
    
    # Format using ChatML
    system_prompt = "You are an expert support ticket analyst. Your task is to extract information from a user's message and format it as a JSON object matching the provided schema."
    
    formatted_prompt = f"<|system|>\n{system_prompt}<|end|>\n<|user|>\nAnalyze the following support ticket and provide the JSON output:\n\nTicket: '{input_text}'<|end|>\n<|assistant|>\n{json.dumps(target_json, indent=2)}<|end|>"
    
    return {"text": formatted_prompt}

# Generate a dataset
num_examples = 500
data = [generate_example() for _ in range(num_examples)]
dataset = Dataset.from_list(data)

# Save the dataset for later use
dataset.to_json("support_ticket_dataset.json")

print(dataset[0]['text'])

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import TrainingArguments
from trl import SFTTrainer

# 1. Prepare model for k-bit training
model.config.use_cache = False # Recommended for training
model = prepare_model_for_kbit_training(model)

# 2. Configure LoRA
lora_config = LoraConfig(
    r=16,  # Rank of the update matrices. A higher rank means more parameters and expressivity.
    lora_alpha=32, # LoRA scaling factor.
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], # Target attention layers
    lora_dropout=0.05, # Dropout for regularization
    bias="none",
    task_type="CAUSAL_LM",
)

# Add LoRA adapter to the model
peft_model = get_peft_model(model, lora_config)
peft_model.print_trainable_parameters() # See how few parameters we are training!

# 3. Configure Training Arguments
training_args = TrainingArguments(
    output_dir="./phi3-lora-json-finetune",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    num_train_epochs=3,
    lr_scheduler_type="cosine",
    save_strategy="epoch",
    logging_steps=10,
    max_steps=-1, # Overridden by num_train_epochs
    fp16=True, # Use mixed precision for speed and memory efficiency
)

# 4. Initialize the Trainer
trainer = SFTTrainer(
    model=peft_model,
    train_dataset=dataset, # Our generated dataset
    dataset_text_field="text", # The column in our dataset containing the formatted text
    max_seq_length=1024, # Adjust based on your context length
    args=training_args,
    peft_config=lora_config,
)

# 5. Start Training
trainer.train()

# Save the trained adapter
trainer.save_model("./phi3-lora-json-adapter")

from peft import PeftModel

# Load the base model (not quantized this time, for merging)
base_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)

# Load the PEFT model with the adapter
peft_model = PeftModel.from_pretrained(base_model, "./phi3-lora-json-adapter")

# Merge the adapter into the base model
merged_model = peft_model.merge_and_unload()

# Save the merged model for easy deployment
merged_model.save_pretrained("./phi3-json-expert-merged")
tokenizer.save_pretrained("./phi3-json-expert-merged")

pip install outlines

import torch
from outlines import models, generate
from pydantic import BaseModel, Field
from typing import Literal

# Load our merged, specialist model
model_path = "./phi3-json-expert-merged"
model = models.transformers(model_path, device="cuda")

# Define the Pydantic schema that matches our training
class SupportTicket(BaseModel):
    username: str
    email: str | None
    category: Literal['login_issue', 'billing_problem', 'feature_request', 'bug_report']
    priority: int = Field(ge=1, le=5)
    summary: str

# The guided generation function from outlines
generate_json = generate.json(model, SupportTicket)

# New user ticket for inference
user_ticket = "Hi, I'm b_smith and my account is locked. I can't access anything. This is a critical issue for my team. My email is [email protected]"

# Create the same prompt structure as in training
system_prompt = "You are an expert support ticket analyst. Your task is to extract information from a user's message and format it as a JSON object matching the provided schema."
prompt = f"<|system|>\n{system_prompt}<|end|>\n<|user|>\nAnalyze the following support ticket and provide the JSON output:\n\nTicket: '{user_ticket}'<|end|>\n<|assistant|>\n"

# Run guided generation
result = generate_json(prompt)

print(result)
# Output will be a Pydantic object, guaranteed to match the schema
# username='b_smith' email='[email protected]' category='login_issue' priority=5 summary="..."