dev

app.py

@@ -83,11 +83,41 @@ 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 = 512
-        #max_length = model.get_input_embeddings().weight.shape[0]
+        def tokenize_function(examples):
+            
+            # Assuming 'text' is the input and 'target' is the expected output
+            model_inputs = tokenizer(
+                examples['text'], 
+                max_length=max_length,  # Set to None for dynamic padding
+                truncation=True,
+                padding='max_length',
+                return_tensors='pt', 
+            )
+        
+            # Setup the decoder input IDs (shifted right)
+            labels = tokenizer(
+                examples['target'], 
+                max_length=max_length,  # Set to None for dynamic padding
+                truncation=True,
+                padding='max_length',
+                text_target=examples['target'],
+                return_tensors='pt', 
+            )
+        
+            # Add labels to the model inputs
+            model_inputs["labels"] = labels["input_ids"]
+            return model_inputs
+    
+        #max_length = 512
+        train_size = len(dataset['train'])
+        half_size = train_size // 2  
+        max_length = model.get_input_embeddings().weight.shape[0]
         try:
-            tokenized_train_dataset = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
-            tokenized_test_dataset = load_from_disk(f'/data/{hub_id.strip()}_test_dataset')
+            tokenized_first_half = load_from_disk(f'/data/{hub_id.strip()}_train_dataset')
+            second_half = dataset['train'].select(range(half_size, train_size))
+            tokenized_second_half = tokenize_function(second_half)
+            tokenized_train_dataset = concatenate_datasets([tokenized_first_half, tokenized_second_half])
+            tokenized_test_dataset = tokenize_function(dataset['test'])
             
             # Create Trainer
             trainer = Trainer(
@@ -99,54 +129,17 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             )            
         except:
             # Load the dataset
-            dataset = load_dataset(dataset_name.strip())
+            dataset = load_dataset(dataset_name.strip())                                  
             tokenizer = AutoTokenizer.from_pretrained('google/t5-efficient-tiny-nh8')
             # Tokenize the dataset
-            def tokenize_function(examples):
-                
-                # Assuming 'text' is the input and 'target' is the expected output
-                model_inputs = tokenizer(
-                    examples['text'], 
-                    max_length=max_length,  # Set to None for dynamic padding
-                    truncation=True,
-                    padding='max_length',
-                    return_tensors='pt', 
-                )
-            
-                # Setup the decoder input IDs (shifted right)
-                labels = tokenizer(
-                    examples['target'], 
-                    max_length=max_length,  # Set to None for dynamic padding
-                    truncation=True,
-                    padding='max_length',
-                    text_target=examples['target'],
-                    return_tensors='pt', 
-                )
             
-                # Add labels to the model inputs
-                model_inputs["labels"] = labels["input_ids"]
-                return model_inputs
         
-            tokenized_datasets = dataset.map(tokenize_function, batched=True)
+            first_half = dataset['train'].select(range(half_size))
+            tokenized_half = tokenize_function(first_half)
             
-            tokenized_datasets['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset')
-            tokenized_datasets['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
-        
-            embedding_size = model.get_input_embeddings().weight.shape[0]
-                        
-            if len(tokenizer) > embedding_size:
-                model.resize_token_embeddings(len(tokenizer))
-                model.resize_position_embeddings(len(tokenizer))
+            tokenized_half.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
         
-            # Create Trainer
-            trainer = Trainer(
-                model=model,
-                args=training_args,
-                train_dataset=tokenized_datasets['train'],
-                eval_dataset=tokenized_datasets['test'],
-                compute_metrics=compute_metrics,
-                #callbacks=[LoggingCallback()], 
-            )            
+            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):