dev

app.py

@@ -86,7 +86,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
     
         tokenizer = AutoTokenizer.from_pretrained('google/t5-efficient-tiny-nh8')
     
-        max_length = model.get_input_embeddings().weight.shape[0]
+        #max_length = model.get_input_embeddings().weight.shape[0]
+        max_length = 512
     
         def tokenize_function(examples):
             
@@ -95,7 +96,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 examples['text'], 
                 max_length=max_length,  # Set to None for dynamic padding
                 truncation=True,
-                padding=True,
+                padding='max_length',
                 return_tensors='pt', 
             )
         
@@ -104,7 +105,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 examples['target'], 
                 max_length=max_length,  # Set to None for dynamic padding
                 truncation=True,
-                padding=True,
+                padding='max_length',
                 #text_target=examples['target'],
                 return_tensors='pt', 
             )
@@ -124,12 +125,14 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 saved_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
                 print("FOUND TEST")                    
                 # Create Trainer
+                data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
                 trainer = Trainer(
                     model=model,
                     args=training_args,
                     train_dataset=train_dataset,
                     eval_dataset=saved_test_dataset,
                     compute_metrics=compute_metrics,
+                    data_collator=data_collator, 
                 )
             
             elif os.access(f'/data/{hub_id.strip()}_train_dataset3', os.R_OK):