ArchCoder/llm-excel-plotter-agent

1import pandas as pd2import torch3from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments4from sklearn.model_selection import train_test_split5 6data = pd.read_csv('data/train_data.csv')7queries = data['query'].tolist()8arguments = data['arguments'].tolist()9 10train_queries, eval_queries, train_arguments, eval_arguments = train_test_split(queries, arguments, test_size=0.2, random_state=42)11 12tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")13model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large")14 15train_encodings = tokenizer(train_queries, truncation=True, padding=True)16eval_encodings = tokenizer(eval_queries, truncation=True, padding=True)17 18with tokenizer.as_target_tokenizer():19    train_labels = tokenizer(train_arguments, truncation=True, padding=True)20    eval_labels = tokenizer(eval_arguments, truncation=True, padding=True)21 22class PlotDataset(torch.utils.data.Dataset):23    def __init__(self, encodings, labels):24        self.encodings = encodings25        self.labels = labels26 27    def __getitem__(self, idx):28        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}29        item['labels'] = torch.tensor(self.labels['input_ids'][idx])30        return item31 32    def __len__(self):33        return len(self.encodings.input_ids)34 35train_dataset = PlotDataset(train_encodings, train_labels)36eval_dataset = PlotDataset(eval_encodings, eval_labels)37 38training_args = Seq2SeqTrainingArguments(39    output_dir='./results',40    per_device_train_batch_size=2,41    per_device_eval_batch_size=2,42    num_train_epochs=3,43    logging_dir='./logs',44    logging_steps=10,45    save_steps=500,46    save_total_limit=2,47    evaluation_strategy="epoch",48    predict_with_generate=True,49    generation_max_length=100,  50)51trainer = Seq2SeqTrainer(52    model=model,53    args=training_args,54    train_dataset=train_dataset,55    eval_dataset=eval_dataset,56    tokenizer=tokenizer,57)58 59trainer.train()60 61trainer.save_model("fine-tuned-bart-large")62tokenizer.save_pretrained("fine-tuned-bart-large")63 64print("Model and tokenizer fine-tuned and saved as 'fine-tuned-bart-large'")