dev

app.py

@@ -2,7 +2,7 @@ import spaces
 import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import DataCollatorForSeq2Seq, AutoConfig
-from datasets import load_dataset, concatenate_datasets, load_from_disk
+from datasets import load_dataset, concatenate_datasets, load_from_disk, DatasetDict
 import traceback
 from sklearn.metrics import accuracy_score
 import numpy as np
@@ -119,7 +119,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             second_half = dataset['train'].select(range(half_size, train_size))
             dataset['train'] = second_half
             tokenized_second_half = dataset.map(tokenize_function, batched=True)
-            tokenized_train_dataset = concatenate_datasets([tokenized_first_half, tokenized_second_half])
+            dataset['train'] = concatenate_datasets([tokenized_first_half['train'], tokenized_second_half['train']])
+            tokenized_train_dataset = dataset['train']
             tokenized_test_dataset = tokenize_function(dataset['test'])
             
             # Create Trainer