dev

app.py

@@ -83,7 +83,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             print("Loading model from checkpoint...")
             model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)        
     
-        max_length = 128
+        max_length = 512
         #max_length = model.get_input_embeddings().weight.shape[0]
         try:
             tokenized_train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
@@ -109,7 +109,8 @@ 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', 
                 )
             
                 # Setup the decoder input IDs (shifted right)
@@ -117,8 +118,9 @@ 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, 
-                    text_target=examples['target']  # Use text_target for target text
+                    padding='max_length',
+                    text_target=examples['target'],
+                    return_tensors='pt', 
                 )
             
                 # Add labels to the model inputs
@@ -178,7 +180,6 @@ def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     config = AutoConfig.from_pretrained("google/t5-efficient-tiny")
     model = AutoModelForSeq2SeqLM.from_config(config)
     initialize_weights(model)
-    print(list(model.named_parameters()))
     lora_config = LoraConfig(
         r=16,  # Rank of the low-rank adaptation
         lora_alpha=32,  # Scaling factor