Fine-tuning is a technique used in machine learning, particularly in natural language processing (NLP), where a pre-trained model is further trained on a specific task or dataset. Fine-tuning a pre-trained model is typically faster than training a model from scratch. This is because the model has already learned valuable features and representations from a large and diverse dataset during pre-training. Fine-tuning often leads to better performance on specific tasks, for example, you can take a pre-trained model trained on a general text corpus and fine-tune it for a domain-specific task, such as medical text analysis or legal document classification. Furthermore, you can adjust the model's behavior and adapt it to the specific needs of your target task. It also helps with costs and prompt limit problem: Once a model has been fine-tuned by training on many more examples than can fit in the prompt, you won't need to provide as many examples in the prompt. This saves costs and enables lower-latency requests.
However, it's important to note that fine-tuning also comes with its own set of limitations and considerations. In practice, fine-tuning requires training a model, with the key distinction being that the weights already contain valuable information. Consequently, fine-tuning shares some of the same needs and challenges as training from scratch. These include the need for task-specific labeled data, the possibility of biases in the pre-trained model, and the risk of overfitting when dealing with small target datasets. Additionally, there are associated costs related to processing requirements. Most large language models (LLMs) demand significant GPU resources for fine-tuning. However, through techniques like Quantization and LoRA and Python code, it is possible to reduce these requirements to the point where fine-tuning the LLama2 7b parameter model can be accomplished with just a single GPU. For more information about this techniques, check this other article.
Code example
Here is a code example demonstrating how to fine-tune the Llama2-7b-chat model in Python. We will utilize the BitsAndBytes library for quantization, PEFT for LoRA to fine-tune specific model parameters, and Transformers to load and train the model. With the help of quantization and LoRA, we were able to train the model using only an ml.g5.xlarge instance of AWS SageMaker, which has only 1 GPU and 24GB of GPU Memory (for the full model we would normally need several GPUs with at least 112 GB of memory in total).
The goal of this fine-tuning process is to enable the model to answer questions about the company Rootstrap without requiring additional context in the prompt. To achieve this, we incorporate Rootstrap-related questions and answers into the training dataset. Additionally, we include unrelated questions in the dataset, all with the same answer: "I am only limited to answer Rootstrap information." This approach ensures that the model specializes in responding to Rootstrap-related queries. To perform the fine tuneLet's go through the code step by step:
Import dependencies
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))Load the Rootstrap’s info dataset
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))Let’s take a look at the 5 firsts elements of the dataset
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))Create bitsandbytes config for 4bits quantization and load the model with the tokenizer that will transform the input text to numeric form:
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))Let’s try the vanilla model without fine-tuning, we will create an inference function for that:
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))
Surprisignly, Llama2 already knows what Rootstrap is, but do not know the projects.
It’s time for fine tuning! First, preprocess the dataset.
Creation of the preprocessing functions:
- Formatting the dataset with the right prompt structure
- Shuffle and removal of unused columns
- Tokenization of the texts inputs to convert to numeric inputs.
- Filter out inputs which are bigger than max input allowed
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))Apply the functions to the dataset:
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))We have our dataset prepared. Now, let’s create the function that will help us with the LoRA configuration, we will use peft library for that.
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))We can finally create and execute our training code:
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))The model was saved in the output_dir.
We can now load it and make it questions about Rootstrap to see if it has improved with fine-tuning:
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))
Not ideal (Rootstrap has never worked with Uber), but it has learnt something indeed. Maybe with more training, or with less quantization, better results can be achieved.
Let’s verify that the model refuses to answer questions not related to Rootstrap:
/* apply_preprocessing.py */
## Preprocess dataset
max_length = get_max_length(model)
dataset = preprocess_dataset(tokenizer, max_length, 123, dataset)
/* dataset.json */
{"question": [
"What does Rootstrap's service of Accessibility & ADA Compliance entail?",
"What role do illustrations play in Rootstrap's services?",
"How does Rootstrap approach UI Visuals?",
"What is the purpose of UX Research at Rootstrap?",
"How does Rootstrap ensure Automation and Top Notch code Quality?"
],
"answer": [
"Rootstrap's Accessibility & ADA Compliance service focuses on creating digital solutions that are accessible to everyone, regardless of age or impairments. The goal is to provide a pleasant experience for all users.",
"Illustrations are used by Rootstrap to attract attention, aid retention, enhance understanding, and create contexts. They aim to boost the product's general look and feel with stunning and unique creations.",
"Rootstrap's experienced team works with clients to provide the most intuitive and up-to-date UI, following the latest visual tendencies in the market. The process begins with moodboards to define the desired style of the product.",
"Rootstrap uses UX research methods to understand the needs of users and expose design problems and opportunities. The insights gained help deliver entirely user-centered solutions, focusing on crafting accessible and engaging user experiences.",
"Rootstrap prides itself on having top-tier code quality and processes. They follow industry-leading standards, code reviews, and test automation to ensure the quality of their deliverables."
]}
/* dependencies.py */
import os
import torch
import bitsandbytes as bnb
from datasets import load_dataset
from functools import partial
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed, Trainer, TrainingArguments, BitsAndBytesConfig, \
DataCollatorForLanguageModeling, Trainer, TrainingArguments
/* fine_tuned_inference.py */
def build_prompt_inference(question):
prompt = "You are a helpful assistant at Rootstrap. If the question is not related to Rootstrap do not answer. ### Instruction:\n\n"
prompt = prompt + question + "\n\n###Response:\n\n"
return prompt
finetuned_model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16, local_files_only=True)
inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), finetuned_model, tokenizer)
/* fine_tuned_inference_2.py */
inference(build_prompt_inference("How to hack the NASA?"), finetuned_model, tokenizer)
/* load_dataset.py */
# Load the databricks dataset from Hugging Face
dataset = load_dataset("csv", data_files='data/rs_info_qa.csv', split="train")
/* load_model.py */
def create_bnb_config():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
return bnb_config
def load_model(model_name, bnb_config):
n_gpus = torch.cuda.device_count()
max_memory = f'{20480}MB'
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto", # dispatch efficiently the model on the available ressources
max_memory = {i: max_memory for i in range(n_gpus)},
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
# Needed for LLaMA tokenizer
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
# Load model from HF with user's token and with bitsandbytes config
model_name = "meta-llama/Llama-2-7b-chat-hf"
bnb_config = create_bnb_config()
model, tokenizer = load_model(model_name, bnb_config)
/* lora_config.py */
def create_peft_config(modules):
"""
Create Parameter-Efficient Fine-Tuning config for your model
:param modules: Names of the modules to apply Lora to
"""
config = LoraConfig(
r=16, # dimension of the updated matrices
lora_alpha=64, # parameter for scaling
target_modules=modules,
lora_dropout=0.1, # dropout probability for layers
bias="none",
task_type="CAUSAL_LM",
)
return config
# SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
def print_trainable_parameters(model, use_4bit=False):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
trainable_params += num_params
if use_4bit:
trainable_params /= 2
print(
f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}"
)
def find_all_linear_names(model):
cls = bnb.nn.Linear4bit #if args.bits == 4 else (bnb.nn.Linear8bitLt if args.bits == 8 else torch.nn.Linear)
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: # needed for 16-bit
lora_module_names.remove('lm_head')
return list(lora_module_names)
/* preprocessing.py */
def create_prompt_formats(sample):
"""
Format various fields of the sample ('instruction', 'context', 'response')
Then concatenate them using two newline characters
:param sample: Sample dictionnary
"""
INTRO_BLURB = "You are a helpfull assistant of a company called Rootstrap. You must answer questions only related to rootstrap, if the question is not related to Rootstrap do not answer."
INSTRUCTION_KEY = "### Question:"
#INPUT_KEY = "Input:"
RESPONSE_KEY = "### Answer:"
END_KEY = "### End"
blurb = f"{INTRO_BLURB}"
instruction = f"{INSTRUCTION_KEY}\n{sample['question']}"
#input_context = f"{INPUT_KEY}\n{sample['context']}" if sample["context"] else None
response = f"{RESPONSE_KEY}\n{sample['answer']}"
end = f"{END_KEY}"
parts = [part for part in [blurb, instruction, response, end] if part]
formatted_prompt = "\n\n".join(parts)
sample["text"] = formatted_prompt
return sample
def get_max_length(model):
conf = model.config
max_length = None
for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]:
max_length = getattr(model.config, length_setting, None)
if max_length:
print(f"Found max lenth: {max_length}")
break
if not max_length:
max_length = 1024
print(f"Using default max length: {max_length}")
return max_length
def preprocess_batch(batch, tokenizer, max_length):
"""
Tokenizing a batch
"""
return tokenizer(
batch["text"],
max_length=max_length,
truncation=True,
)
# SOURCE https://github.com/databrickslabs/dolly/blob/master/training/trainer.py
def preprocess_dataset(tokenizer: AutoTokenizer, max_length: int, seed, dataset: str):
"""Format & tokenize it so it is ready for training
:param tokenizer (AutoTokenizer): Model Tokenizer
:param max_length (int): Maximum number of tokens to emit from tokenizer
"""
# Add prompt to each sample
print("Preprocessing dataset...")
dataset = dataset.map(create_prompt_formats)#, batched=True)
print(dataset[0])
# Apply preprocessing to each batch of the dataset & and remove unused fields
_preprocessing_function = partial(preprocess_batch, max_length=max_length, tokenizer=tokenizer)
dataset = dataset.map(
_preprocessing_function,
batched=True,
remove_columns=["question", "text", "answer"],
)
# Filter out samples that have input_ids exceeding max_length
dataset = dataset.filter(lambda sample: len(sample["input_ids"]) < max_length)
# Shuffle dataset
dataset = dataset.shuffle(seed=seed)
return dataset
/* print_dataset.py */
print(dataset[:5])
/* train.py */
def train(model, tokenizer, dataset, output_dir):
# Apply preprocessing to the model to prepare it by
# 1 - Enabling gradient checkpointing to reduce memory usage during fine-tuning
model.gradient_checkpointing_enable()
# 2 - Using the prepare_model_for_kbit_training method from PEFT
model = prepare_model_for_kbit_training(model)
# Get lora module names
modules = find_all_linear_names(model)
# Create PEFT config for these modules and wrap the model to PEFT
peft_config = create_peft_config(modules)
model = get_peft_model(model, peft_config)
# Print information about the percentage of trainable parameters
print_trainable_parameters(model)
# Training parameters
trainer = Trainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=25,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
optim="paged_adamw_8bit",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
model.config.use_cache = False # re-enable for inference to speed up predictions for similar inputs
### SOURCE https://github.com/artidoro/qlora/blob/main/qlora.py
# Verifying the datatypes before training
dtypes = {}
for _, p in model.named_parameters():
dtype = p.dtype
if dtype not in dtypes: dtypes[dtype] = 0
dtypes[dtype] += p.numel()
total = 0
for k, v in dtypes.items(): total+= v
for k, v in dtypes.items():
print(k, v, v/total)
do_train = True
# Launch training
print("Training...")
if do_train:
train_result = trainer.train()
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
print(metrics)
###
# Saving model
print("Saving last checkpoint of the model...")
os.makedirs(output_dir, exist_ok=True)
trainer.save_model(output_dir)
# Free memory for merging weights
del model
del trainer
torch.cuda.empty_cache()
output_dir = "results/llama2/final_checkpoint"
train(model, tokenizer, dataset, output_dir)
/* vanilla_inference.py */
def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=100):
# Tokenize
input_ids = tokenizer.encode(
text,
return_tensors="pt",
truncation=True,
max_length=max_input_tokens
)
# Generate
device = model.device
generated_tokens_with_prompt = model.generate(
input_ids=input_ids.to(device),
max_length=max_output_tokens,
temperature=0.001
)
# Decode
generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
# Strip the prompt
generated_text_answer = generated_text_with_prompt[0][len(text):]
return generated_text_answer
print(inference(build_prompt_inference("Can you tell me some projects carried out by rootstrap?"), model, tokenizer))A little verbose, but refuses correctly to answer the question.
Now we can use this fine-tuned model, for example, to create a simple python Q&A chatbot that answers Rootstrap’s information. Although it still may hallucinate a little without context, we don't have to pass all the company context in every question, since it already has some of the information embedded and has learned not to answer questions unrelated to Rootstrap.
