dev

app.py

@@ -109,20 +109,22 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 #text_target=examples['target'],
                 return_tensors='pt', 
             )
-        
+            labels["input_ids"] = [
+                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
+            ]        
             # Add labels to the model inputs
             model_inputs["labels"] = labels["input_ids"]
             return model_inputs
     
         #max_length = 512
         # Load the dataset        
-        
+        column_names = ['text', 'target']
         
         try:
             saved_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
             if os.access(f'/data/{hub_id.strip()}_test_dataset', os.R_OK):
                 train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset3')
-                saved_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_validation_dataset')
+                saved_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
                 print("FOUND TEST")                    
                 # Create Trainer
                 data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
@@ -132,14 +134,15 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                     train_dataset=train_dataset,
                     eval_dataset=saved_test_dataset,
                     compute_metrics=compute_metrics,
-                    data_collator=data_collator, 
+                    data_collator=data_collator,
+                    processing_class=tokenizer,
                 )
             
             elif os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):                        
                 dataset = load_dataset(dataset_name.strip())
                 del dataset['train']
                 del dataset['validation']
-                test_set = dataset.map(tokenize_function, batched=True, batch_size=50)
+                test_set = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
                 test_set['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
                 return 'TRAINING DONE'
                 
@@ -154,7 +157,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 third_third = dataset['train'].select(range(third_size*2, train_size))
                 dataset['train'] = third_third
                 #tokenized_second_half = tokenize_function(third_third)
-                tokenized_second_half = dataset.map(tokenize_function, batched=True, batch_size=50)
+                tokenized_second_half = dataset.map(tokenize_function, batched=True, batch_size=50,remove_columns=column_names,)
                 dataset['train'] = concatenate_datasets([saved_dataset, tokenized_second_half['train']])
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset3')
                 return 'THIRD THIRD LOADED'
@@ -167,7 +170,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 second_third = dataset['train'].select(range(third_size, third_size*2))
                 dataset['train'] = second_third
                 del dataset['test']
-                tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True, batch_size=50)
+                tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
                 dataset['train'] = concatenate_datasets([saved_dataset['train'], tokenized_sh_fq_dataset['train']])
                 dataset['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset2')
                 dataset['validation'].save_to_disk(f'/data/{hub_id.strip()}_validation_dataset')
@@ -183,7 +186,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             dataset['train'] = first_third
             del dataset['test']
             del dataset['validation']
-            tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=50)
+            tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=50, remove_columns=column_names,)
             
             tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
             print('DONE')