init

app.py

@@ -3,11 +3,13 @@ import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
 from datasets import load_dataset
 import traceback
+from huggingface_hub import login
 from peft import get_peft_model, LoraConfig
 
 @spaces.GPU
-def fine_tune_model(model_name, dataset_name, hub_id, num_epochs, batch_size, lr, grad):
+def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     try:
+        login(api_key)
         lora_config = LoraConfig(
             r=16,  # Rank of the low-rank adaptation
             lora_alpha=32,  # Scaling factor
@@ -25,7 +27,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, num_epochs, batch_size, lr
         # Tokenize the dataset
         def tokenize_function(examples):
             max_length = 256
-            return tokenizer(examples['text'], truncation=True)
+            return tokenizer(examples['text'], max_length=max_length, truncation=True)
     
         tokenized_datasets = dataset.map(tokenize_function, batched=True)
     
@@ -82,6 +84,7 @@ try:
             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"),
             gr.Slider(minimum=1, maximum=16, value=4, label="Batch Size"),
             gr.Slider(minimum=1, maximum=1000, value=50, label="Learning Rate (e-5)"),

requirements.txt

@@ -2,3 +2,4 @@ spaces
 transformers
 datasets
 peft
+huggingface_hub