init

app.py

@@ -1,11 +1,24 @@
 import spaces
 import gradio as gr
-from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM, TrainerCallback
 from datasets import load_dataset
 import traceback
 from huggingface_hub import login
 from peft import get_peft_model, LoraConfig
 
+class LoggingCallback(TrainerCallback):
+    def on_step_end(self, args, state, control, kwargs):
+        # Log the learning rate
+        current_lr = state.optimizer.param_groups[0]['lr']
+        print(f"Current Learning Rate: {current_lr}")
+
+    def on_epoch_end(self, args, state, control, kwargs):
+        # Log the error rate (assuming you have a metric to calculate it)
+        # Here we assume you have a way to get the validation loss
+        if state.best_metric is not None:
+            error_rate = 1 - state.best_metric  # Assuming best_metric is accuracy
+            print(f"Current Error Rate: {error_rate:.4f}")
+
 @spaces.GPU
 def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     try:
@@ -21,7 +34,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
     
         # Load the model and tokenizer
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
-        #model = get_peft_model(model, lora_config)
+        model = get_peft_model(model, lora_config)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
     
         # Tokenize the dataset
@@ -37,9 +50,9 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             eval_strategy="epoch",
             save_strategy='epoch', 
             learning_rate=lr*0.00001,
-            per_device_train_batch_size=batch_size,
-            per_device_eval_batch_size=batch_size, 
-            num_train_epochs=num_epochs,
+            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,
             gradient_accumulation_steps=grad*0.1,
             load_best_model_at_end=True,
@@ -59,6 +72,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             args=training_args,
             train_dataset=tokenized_datasets['train'],
             eval_dataset=tokenized_datasets['validation'],
+            callbacks=[LoggingCallback()], 
         )
     
         # Fine-tune the model