dev

app.py

@@ -55,6 +55,8 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                  
         # Set training arguments
         training_args = TrainingArguments(
+            torch_empty_cache_steps=150, 
+            overwrite_output_dir=True, 
             output_dir='/data/results',
             eval_strategy="steps",  # Change this to steps
             save_strategy='steps',
@@ -69,18 +71,18 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             metric_for_best_model="accuracy",
             greater_is_better=True,
             logging_dir='/data/logs',
-            logging_steps=10,
+            logging_steps=250,
             #push_to_hub=True,
             hub_model_id=hub_id.strip(),
             fp16=True,
             #lr_scheduler_type='cosine',
-            save_steps=500,  # Save checkpoint every 500 steps
+            save_steps=350,  # Save checkpoint every 500 steps
             save_total_limit=3, 
         )
         # Check if a checkpoint exists and load it
-        #if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
-            #print("Loading model from checkpoint...")
-            #model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)        
+        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
+            print("Loading model from checkpoint...")
+            model = AutoModelForSeq2SeqLM.from_pretrained(training_args.output_dir)        
     
         tokenizer = AutoTokenizer.from_pretrained('google/t5-efficient-tiny-nh8')
     
@@ -136,7 +138,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
                 dataset['train'] = dataset['train'].select(range(8000))
                 del dataset['train']
                 del dataset['validation']
-                test_set = dataset.map(tokenize_function, batched=True)
+                test_set = dataset.map(tokenize_function, batched=True, batch_size=20)
                 test_set['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
                 return 'TRAINING DONE'
                 
@@ -152,7 +154,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)
+                tokenized_second_half = dataset.map(tokenize_function, batched=True, batch_size=20)
                 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 +169,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)
+                tokenized_sh_fq_dataset = dataset.map(tokenize_function, batched=True, batch_size=20)
                 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')
@@ -184,15 +186,15 @@ 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)
+            tokenized_first_third = dataset.map(tokenize_function, batched=True, batch_size=20)
             
             tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
             print('DONE')
             return 'RUN AGAIN TO LOAD REST OF DATA'
 
         # Fine-tune the model
-        #if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
-            #train_result = trainer.train(resume_from_checkpoint=True)
+        if os.path.exists(training_args.output_dir) and os.listdir(training_args.output_dir):
+            train_result = trainer.train(resume_from_checkpoint=True)
         else:
             train_result = trainer.train()
         trainer.push_to_hub(commit_message="Training complete!")