init

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM, TrainerCallback
 from transformers import DataCollatorForSeq2Seq
 from datasets import load_dataset
+from datasets import concatenate_datasets
 import traceback
 from huggingface_hub import login
 from peft import get_peft_model, LoraConfig
@@ -43,30 +44,43 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
     
         # Tokenize the dataset
         def tokenize_function(examples):
-            
-            # Assuming 'text' is the input and 'target' is the expected output
+            # Tokenize the input text
             model_inputs = tokenizer(
                 examples['text'], 
-                max_length=max_length,  # Set to None for dynamic padding
-                padding=False,     # Disable padding here, we will handle it later
+                max_length=max_length,  
+                padding=True,     
                 truncation=True,
             )
         
-            # Setup the decoder input IDs (shifted right)
+            # Tokenize the target text
             labels = tokenizer(
                 examples['target'], 
-                max_length=max_length,  # Set to None for dynamic padding
-                padding=False,     # Disable padding here, we will handle it later
+                max_length=max_length,  
+                padding=True,     
                 truncation=True,
-                text_target=examples['target']  # Use text_target for target text
             )
         
             # Add labels to the model inputs
             model_inputs["labels"] = labels["input_ids"]
             return model_inputs
-    
-        tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=1)
-        data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+        
+        # Define a function to process the dataset in chunks
+        def tokenize_in_chunks(dataset, chunk_size=1000):
+            tokenized_datasets = []
+            
+            for i in range(0, len(dataset), chunk_size):
+                chunk = dataset[i:i + chunk_size]
+                tokenized_chunk = chunk.map(tokenize_function, batched=True)
+                tokenized_datasets.append(tokenized_chunk)
+            
+            # Concatenate all tokenized chunks into a single dataset
+            return tokenized_datasets
+        
+        # Tokenize the dataset in chunks
+        tokenized_datasets = tokenize_in_chunks(dataset['train'], chunk_size=1000)
+        
+        # If you want to combine all chunks into a single dataset
+        final_tokenized_dataset = concatenate_datasets(tokenized_datasets)
     
         # Set training arguments
         training_args = TrainingArguments(
@@ -94,8 +108,8 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=data_collator['train'],
-            eval_dataset=data_collator['test'],
+            train_dataset=final_tokenized_dataset['train'],
+            eval_dataset=final_tokenized_dataset['test'],
             #callbacks=[LoggingCallback()], 
         )