Upload 2 files

app.py

@@ -1,88 +1,137 @@
+import spaces
 import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
-from datasets import load_dataset
+from transformers import DataCollatorForSeq2Seq
+from datasets import load_dataset, concatenate_datasets, load_from_disk
 import traceback
 
-def fine_tune_model(model_name, dataset_name, hub_id, num_epochs, batch_size, lr, grad):
-    '''
+
+import os
+from huggingface_hub import login
+
+@spaces.GPU
+def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     try:
-        
-        # Load the dataset
-        dataset = load_dataset(dataset_name)
-    
+        #login(api_key.strip())
         # Load the model and tokenizer
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, num_labels=2)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-        # Tokenize the dataset
-        def tokenize_function(examples):
-            return tokenizer(examples['text'], padding="max_length", truncation=True)
-    
-        tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
+        
     
         # Set training arguments
         training_args = TrainingArguments(
-            output_dir='./results',
-            evaluation_strategy="epoch",
-            learning_rate=lr,
-            per_device_train_batch_size=batch_size,
-            per_device_eval_batch_size=batch_size, 
-            num_train_epochs=num_epochs,
+            output_dir='/data/results',
+            eval_strategy="steps",  # Change this to steps
+            save_strategy='steps',
+            learning_rate=lr*0.00001,
+            per_device_train_batch_size=int(batch_size),
+            per_device_eval_batch_size=int(batch_size), 
+            num_train_epochs=int(num_epochs),
             weight_decay=0.01,
-            evaluation_strategy='epoch',
-            gradient_accumulation_steps=grad,
+            gradient_accumulation_steps=int(grad),
+            max_grad_norm = 1.0, 
             load_best_model_at_end=True,
             metric_for_best_model="accuracy",
             greater_is_better=True,
-            logging_dir='./logs',
+            logging_dir='/data/logs',
             logging_steps=10,
-            push_to_hub=True,
-            hub_model_id=hub_id, 
+            #push_to_hub=True,
+            hub_model_id=hub_id.strip(),
+            fp16=True,
+            #lr_scheduler_type='cosine',
+            save_steps=100,  # Save checkpoint every 500 steps
+            save_total_limit=3, 
         )
+        # Check if a checkpoint exists and load it
+        max_length = 128
+        # Load the dataset
+        dataset = load_dataset(dataset_name.strip())
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        # 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
+            )
+        
+            # Add labels to the model inputs
+            model_inputs["labels"] = labels["input_ids"]
+        
+        
+        tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        
+        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')
     
         # Create Trainer
         trainer = Trainer(
             model=model,
             args=training_args,
             train_dataset=tokenized_datasets['train'],
-            eval_dataset=tokenized_datasets['validation'],
-        )
-    
+            eval_dataset=tokenized_datasets['test'],
+            compute_metrics=compute_metrics,
+            #callbacks=[LoggingCallback()], 
+        )            
+
         # Fine-tune the model
         trainer.train()
         trainer.push_to_hub(commit_message="Training complete!")
     except Exception as e:
         return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"
-    '''
     return 'DONE!'#model
 '''
 # Define Gradio interface
 def predict(text):
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
     outputs = model(inputs)
     predictions = outputs.logits.argmax(dim=-1)
-    return "Positive" if predictions.item() == 1 else "Negative"
+    return predictions.item()
 '''
 # Create Gradio interface
 try:
-    
     iface = gr.Interface(
-        fine_tune_model,
+        fn=fine_tune_model,
         inputs=[
-            gr.inputs.Textbox(label="Model Name (e.g., 'google/t5-efficient-tiny-nh8')"),
-            gr.inputs.Textbox(label="Dataset Name (e.g., 'imdb')"),
-            gr.inputs.Textbox(label="HF hub to push to after training"),
-            gr.inputs.Slider(minimum=1, maximum=10, default=3, label="Number of Epochs"),
-            gr.inputs.Slider(minimum=1, maximum=16, default=4, label="Batch Size"),
-            gr.inputs.Slider(minimum=1, maximum=100, default=50, label="Learning Rate (e-5)"),
-            gr.inputs.Slider(minimum=1, maximum=100, default=1, label="Gradient accumulation (e-1)"), 
+            gr.Textbox(label="Model Name (e.g., 'google/t5-efficient-tiny-nh8')"),
+            gr.Textbox(label="Dataset Name (e.g., 'imdb')"),
+            gr.Textbox(label="HF hub to push to after training"),
+            gr.Textbox(label="HF API token"),
+            gr.Slider(minimum=1, maximum=10, value=3, label="Number of Epochs", step=1),
+            gr.Slider(minimum=1, maximum=2000, value=1, label="Batch Size", step=1),
+            gr.Slider(minimum=1, maximum=1000, value=1, label="Learning Rate (e-5)", step=1),
+            gr.Slider(minimum=1, maximum=100, value=1, label="Gradient accumulation", step=1), 
         ],
         outputs="text",
         title="Fine-Tune Hugging Face Model",
         description="This interface allows you to fine-tune a Hugging Face model on a specified dataset."
     )
+    '''
+    iface = gr.Interface(
+        fn=predict,
+        inputs=[
+            gr.Textbox(label="Query"),
+        ],
+        outputs="text",
+        title="Fine-Tune Hugging Face Model",
+        description="This interface allows you to test a fine-tune Hugging Face model."
+    )
+    '''
+    # Launch the interface
+    iface.launch()    
 except Exception as e:
     print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
 
-# Launch the interface
-iface.launch()

requirements.txt

@@ -0,0 +1,8 @@
+spaces
+transformers
+datasets
+peft
+huggingface_hub
+scikit-learn
+numpy
+torch