init

app.py

@@ -3,7 +3,6 @@ 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
@@ -44,43 +43,30 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
     
         # Tokenize the dataset
         def tokenize_function(examples):
-            # Tokenize the input text
+            
+            # Assuming 'text' is the input and 'target' is the expected output
             model_inputs = tokenizer(
                 examples['text'], 
-                max_length=max_length,  
-                padding=True,     
+                max_length=max_length,  # Set to None for dynamic padding
+                padding=False,     # Disable padding here, we will handle it later
                 truncation=True,
             )
         
-            # Tokenize the target text
+            # Setup the decoder input IDs (shifted right)
             labels = tokenizer(
                 examples['target'], 
-                max_length=max_length,  
-                padding=True,     
+                max_length=max_length,  # Set to None for dynamic padding
+                padding=False,     # Disable padding here, we will handle it later
                 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
-        
-        # 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, chunk_size=1000)
-        
-        # If you want to combine all chunks into a single dataset
-        final_tokenized_dataset = concatenate_datasets(tokenized_datasets)
+    
+        tokenized_datasets = dataset.map(tokenize_function)
+        data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
     
         # Set training arguments
         training_args = TrainingArguments(
@@ -108,8 +94,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=final_tokenized_dataset['train'],
-            eval_dataset=final_tokenized_dataset['test'],
+            train_dataset=data_collator['train'],
+            eval_dataset=data_collator['test'],
             #callbacks=[LoggingCallback()], 
         )