Update app.py

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 import os, torch
 from datasets import load_dataset
 from huggingface_hub import HfApi, login
-from peft import LoraConfig, TaskType, get_peft_model
+from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
 from transformers import AutoModelForCausalLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, pipeline
 
 ACTION_1 = "Prompt base model"
@@ -102,8 +102,10 @@ def fine_tune_model(base_model_name, dataset_name):
         #target_modules=["q", "v"],
         task_type=TaskType.SEQ_2_SEQ_LM,
     )
+
+    model = prepare_model_for_kbit_training(model)
     
-    peft_model = get_peft_model(model, lora_config)
+    model = get_peft_model(model, lora_config)
     
     print("### PEFT")
     peft_model.print_trainable_parameters() # trainable params: 6,815,744 || all params: 8,037,076,992 || trainable%: 0.0848