|
import torch |
|
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments |
|
from datasets import load_dataset, Dataset |
|
from peft import LoraConfig, get_peft_model |
|
from transformers import BitsAndBytesConfig, DataCollatorForLanguageModeling |
|
import evaluate |
|
import pandas as pd |
|
import time |
|
|
|
|
|
bnb_config = BitsAndBytesConfig( |
|
load_in_4bit=True, |
|
bnb_4bit_quant_type="nf4", |
|
bnb_4bit_use_double_quant=True, |
|
bnb_4bit_compute_dtype=torch.float16 |
|
) |
|
|
|
|
|
model_name = "meta-llama/Llama-2-7b-chat-hf" |
|
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True) |
|
|
|
|
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_name, |
|
quantization_config=bnb_config, |
|
device_map="auto", |
|
use_auth_token=True |
|
) |
|
|
|
print(f"Model {model_name} loaded and quantized in 4 bits") |
|
|
|
|
|
dataset = load_dataset("andreshere/counsel_chat", use_auth_token=True) |
|
print(f"Dataset loaded") |
|
|
|
|
|
lora_config = LoraConfig( |
|
r=32, |
|
lora_alpha=32, |
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], |
|
lora_dropout=0.1 |
|
) |
|
|
|
model = get_peft_model(model, lora_config) |
|
|
|
|
|
def preprocess_function(examples): |
|
inputs = examples['Context'] |
|
targets = examples['Response'] |
|
model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding='max_length') |
|
labels = tokenizer(targets, max_length=512, truncation=True, padding='max_length').input_ids |
|
|
|
model_inputs["labels"] = labels |
|
model_inputs["Context"] = inputs |
|
model_inputs["Response"] = targets |
|
return model_inputs |
|
|
|
tokenized_dataset = dataset.map(preprocess_function, batched=True) |
|
|
|
train_dataset = tokenized_dataset['train'] |
|
eval_dataset = tokenized_dataset['test'] |
|
|
|
data_collator = DataCollatorForLanguageModeling( |
|
tokenizer=tokenizer, |
|
mlm=False, |
|
) |
|
|
|
|
|
training_args = TrainingArguments( |
|
output_dir='./', |
|
do_train=True, |
|
do_eval=True, |
|
eval_strategy='epoch', |
|
gradient_accumulation_steps=2, |
|
auto_find_batch_size=True, |
|
weight_decay=0.01, |
|
num_train_epochs=4, |
|
learning_rate=1e-4, |
|
logging_dir='./logs', |
|
logging_strategy="steps", |
|
logging_steps=10, |
|
save_strategy="epoch", |
|
save_total_limit=2, |
|
save_safetensors=True, |
|
eval_steps=10, |
|
report_to="tensorboard", |
|
hub_strategy="every_save" |
|
) |
|
|
|
|
|
trainer = Trainer( |
|
model=model, |
|
args=training_args, |
|
train_dataset=train_dataset, |
|
eval_dataset=eval_dataset, |
|
data_collator=data_collator |
|
) |
|
|
|
|
|
trainer.train() |
|
|
|
|
|
trainer.save_model() |
|
|
|
|
|
trainer.push_to_hub("andreshere/llama-2-7b-mental-health-counseler") |
|
|
|
|
|
rouge = evaluate.load('rouge') |
|
bleu = evaluate.load('bleu') |
|
bertscore = evaluate.load('bertscore') |
|
meteor = evaluate.load('meteor') |
|
|
|
|
|
def compute_metrics_and_log_time(pred): |
|
labels_ids = pred.label_ids |
|
pred_ids = pred.predictions |
|
decoded_preds = tokenizer.batch_decode(pred_ids, skip_special_tokens=True) |
|
decoded_labels = tokenizer.batch_decode(labels_ids, skip_special_tokens=True) |
|
|
|
|
|
rouge_result = rouge.compute(predictions=decoded_preds, references=decoded_labels) |
|
bleu_result = bleu.compute(predictions=decoded_preds, references=decoded_labels) |
|
bertscore_result = bertscore.compute(predictions=decoded_preds, references=decoded_labels) |
|
meteor_result = meteor.compute(predictions=decoded_preds, references=decoded_labels) |
|
|
|
result = { |
|
"rouge": rouge_result, |
|
"bleu": bleu_result, |
|
"bertscore": bertscore_result, |
|
"meteor": meteor_result, |
|
} |
|
|
|
|
|
compute_times = [] |
|
contexts = [] |
|
original_responses = [] |
|
model_responses = [] |
|
|
|
for i, context in enumerate(pred.features['Context']): |
|
start_time = time.time() |
|
model_response = decoded_preds[i] |
|
end_time = time.time() |
|
compute_time = end_time - start_time |
|
|
|
compute_times.append(compute_time) |
|
contexts.append(context) |
|
original_responses.append(decoded_labels[i]) |
|
model_responses.append(model_response) |
|
|
|
|
|
log_df = pd.DataFrame({ |
|
"Context": contexts, |
|
"Original Response": original_responses, |
|
"Model Response": model_responses, |
|
"Compute Time": compute_times |
|
}) |
|
|
|
log_df.to_csv("response_log.csv", index=False) |
|
|
|
return result |
|
|
|
|
|
metrics = trainer.evaluate(eval_dataset=eval_dataset, metric_key_prefix="eval", compute_metrics=compute_metrics_and_log_time) |
|
print(metrics) |
|
|
|
|
|
metrics_df = pd.DataFrame(metrics, index=[0]) |
|
metrics_df.to_csv("fine-tuned-metrics.csv", index=False) |
|
|
|
|