init

app.py

@@ -3,16 +3,23 @@ import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
 from datasets import load_dataset
 import traceback
+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):
     try:
-        
+        lora_config = LoraConfig(
+            r=16,  # Rank of the low-rank adaptation
+            lora_alpha=32,  # Scaling factor
+            lora_dropout=0.1,  # Dropout for LoRA layers
+            bias="none"  # Bias handling
+        )        
         # Load the dataset
         dataset = load_dataset(dataset_name)
     
         # Load the model and tokenizer
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name, num_labels=2)
+        model = get_peft_model(model, lora_config)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
     
         # Tokenize the dataset
@@ -37,7 +44,9 @@ def fine_tune_model(model_name, dataset_name, hub_id, num_epochs, batch_size, lr
             logging_dir='./logs',
             logging_steps=10,
             push_to_hub=True,
-            hub_model_id=hub_id, 
+            hub_model_id=hub_id,
+            fp16=True,
+            lr_scheduler_type='cosine',
         )
     
         # Create Trainer