dev

app.py

@@ -116,12 +116,20 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
         max_length = model.get_input_embeddings().weight.shape[0]
         try:
             tokenized_first_half = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
-            second_half = dataset['train'].select(range(half_size, train_size))
-            dataset['train'] = second_half
-            tokenized_second_half = dataset.map(tokenize_function, batched=True)
-            dataset['train'] = concatenate_datasets([tokenized_first_half['train'], tokenized_second_half['train']])
-            tokenized_train_dataset = dataset['train']
-            tokenized_test_dataset = tokenize_function(dataset['test'])
+            if 'test' in tokenized_first_half.keys():                
+                second_half_second_quarter = dataset['train'].select(range(half_size+half_size//2, train_size))
+                dataset['train'] = second_half_second_quarter
+                tokenized_second_half = dataset.map(tokenize_function, batched=True)
+                dataset['train'] = concatenate_datasets([tokenized_first_half['train'], tokenized_second_half['train']])
+                tokenized_train_dataset = dataset['train']
+                tokenized_test_dataset = tokenize_function(dataset['test'])
+            else:
+                second_half_first_quarter = dataset['train'].select(range(half_size, half_size+half_size//2))
+                dataset['train'] = second_half_first_quarter
+                tokenized_second_half = dataset.map(tokenize_function, batched=True)
+                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
             trainer = Trainer(