Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import os, torch
 from datasets import load_dataset
 from huggingface_hub import HfApi, login
+from peft import LoraConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, pipeline
 
 ACTION_1 = "Prompt base model"
@@ -80,8 +81,8 @@ def fine_tune_model(base_model_name, dataset_name):
 
     training_args = Seq2SeqTrainingArguments(
         output_dir=f"./{FT_MODEL_NAME}",
-        num_train_epochs=3,
-        #max_steps=1, # overwrites num_train_epochs
+        num_train_epochs=3, # 37,500 steps
+        max_steps=1, # overwrites num_train_epochs
         push_to_hub=True, # only pushes model, also need to push tokenizer (see below)
         # TODO https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments
     )
@@ -89,6 +90,14 @@ def fine_tune_model(base_model_name, dataset_name):
     print("### Training arguments")
     print(training_args)
     print("###")
+
+    # PEFT
+
+    peft_config = LoraConfig(
+        r=8,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
     
     # Create trainer
 
@@ -97,6 +106,7 @@ def fine_tune_model(base_model_name, dataset_name):
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
+        peft_config=peft_config,
         # TODO https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainer
     )