init

app.py

@@ -2,7 +2,7 @@ import spaces
 import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import DataCollatorForSeq2Seq
-from datasets import load_dataset
+from datasets import load_dataset, concatenate_datasets, load_from_disk
 import traceback
 import os
 from huggingface_hub import login
@@ -27,33 +27,40 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         tokenizer = AutoTokenizer.from_pretrained(model_name)
     
         max_length = 64
-    
-        # Tokenize the dataset
-        def tokenize_function(examples):
+        try:
+            tokenized_train_dataset = load_from_disk(f'{hub_id.strip()}_train_dataset')
+            tokenized_test_dataset = load_from_disk(f'{hub_id.strip()}_test_dataset')
+            tokenized_datasets = concatenate_datasets([tokenized_train_dataset, tokenized_test_dataset])
+        except:
+            # 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
+                    padding=True,     # Disable padding here, we will handle it later
+                    truncation=True,
+                )
+            
+                # Setup the decoder input IDs (shifted right)
+                labels = tokenizer(
+                    examples['target'], 
+                    max_length=max_length,  # Set to None for dynamic padding
+                    padding=True,     # Disable padding here, we will handle it later
+                    truncation=True,
+                    text_target=examples['target']  # Use text_target for target text
+                )
             
-            # 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
-                padding=True,     # Disable padding here, we will handle it later
-                truncation=True,
-            )
+                # Add labels to the model inputs
+                model_inputs["labels"] = labels["input_ids"]
+                return model_inputs
         
-            # Setup the decoder input IDs (shifted right)
-            labels = tokenizer(
-                examples['target'], 
-                max_length=max_length,  # Set to None for dynamic padding
-                padding=True,     # Disable padding here, we will handle it later
-                truncation=True,
-                text_target=examples['target']  # Use text_target for target text
-            )
+            tokenized_datasets = dataset.map(tokenize_function, batched=True)
+            
+            tokenized_datasets['train'].save_to_disk(f'{hub_id.strip()}_train_dataset')
+            tokenized_datasets['validation'].save_to_disk(f'{hub_id.strip()}_test_dataset')
         
-            # Add labels to the model inputs
-            model_inputs["labels"] = labels["input_ids"]
-            return model_inputs
-    
-        tokenized_datasets = dataset.map(tokenize_function, batched=True)
-        data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
     
         # Set training arguments
         training_args = TrainingArguments(