fine_tuning / fine_tuning.py
aziz100's picture
Update fine_tuning.py
b827276 verified
from datasets import load_dataset
from transformers import BartTokenizer, BartForSequenceClassification, Trainer, TrainingArguments
import pandas as pd
from datasets import load_dataset, DatasetDict
dataset = load_dataset("csv", data_files="/home/aziz/fine_tuning/FAQ_Appliance_Store_FR.csv")
split_dataset = dataset["train"].train_test_split(test_size=0.2)
dataset = DatasetDict({
"train": split_dataset["train"],
"test": split_dataset["test"]
})
# Load pretrained model and tokenizer
model = BartForSequenceClassification.from_pretrained("facebook/bart-large-mnli")
tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-mnli")
# Tokenize the dataset
def preprocess_function(examples):
return tokenizer(examples['question'], examples['answer'], truncation=True, padding="max_length")
tokenized_datasets = dataset.map(preprocess_function, batched=True)
# Define training arguments
training_args = TrainingArguments(
output_dir="./results",
evaluation_strategy="epoch",
save_strategy="epoch",
per_device_train_batch_size=8,
num_train_epochs=3,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["test"],
)
trainer.train()
model.save_pretrained("./my_model")
tokenizer.save_pretrained("./my_model")