chimp/src/train.py

import pandas as pd
from transformers import TrainingArguments
from trl import SFTTrainer
from peft import LoraConfig, PeftModel, get_peft_model
from model import load_model
import config # Make sure you have a valid config module
from dataset import CustomDataset, template_dataset
from datasets import Dataset, Features, Value, Sequence


if __name__ == '__main__':

 model, tokenizer, peft_config = load_model(config.model_name)
 
 df = pd.read_csv('../data/trainv2.csv') #data
 data_list = df.to_dict(orient='records')

 # custom dataset object
 custom_dataset = CustomDataset(data_list)
 df = pd.DataFrame(custom_dataset.data, columns=["instruction", "context", "response"])
 # Dataset features
 features = Features({
 "instruction": Value("string"),
 "context": Value("string"),
 "response": Value("string"),
 })

 # Create a Hugging Face Dataset from the Pandas DataFrame
 hugging_face_dataset = Dataset.from_pandas(df, features=features)
 dataset = hugging_face_dataset.map(lambda x: template_dataset(x, tokenizer), remove_columns=list(hugging_face_dataset.features))
 print("----training data structure----",dataset)

 # Training Arguments
 training_arguments = TrainingArguments(
 output_dir=config.output_dir,
 per_device_train_batch_size=config.per_device_train_batch_size,
 gradient_accumulation_steps=config.gradient_accumulation_steps,
 optim=config.optim,
 save_steps=config.save_steps,
 logging_steps=config.logging_steps,
 learning_rate=config.learning_rate,
 fp16=config.fp16,
 bf16=config.bf16,
 max_grad_norm=config.max_grad_norm,
 max_steps=config.max_steps,
 warmup_ratio=config.warmup_ratio,
 group_by_length=config.group_by_length,
 lr_scheduler_type=config.lr_scheduler_type,
 report_to="tensorboard"
 )

 # SFTTrainer
 trainer = SFTTrainer(
 model=model,
 train_dataset=dataset,
 peft_config=peft_config,
 dataset_text_field="text",
 max_seq_length=config.max_seq_length,
 tokenizer=tokenizer,
 args=training_arguments,
 packing=config.packing,
 )
 print("**************** TRAINING STARTED ****************")
 trainer.train()
 trainer.model.save_pretrained(config.output_dir)
 print("**************** TRAINING OVER ****************")