Update app.py

app.py

@@ -5,7 +5,8 @@ import os, torch
 from datasets import load_dataset
 from huggingface_hub import HfApi, login
 #from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
-from transformers import AutoModelForCausalLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, pipeline
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, Trainer, TrainingArguments, pipeline
 
 ACTION_1 = "Prompt base model"
 ACTION_2 = "Fine-tune base model"
@@ -79,7 +80,7 @@ def fine_tune_model(base_model_name, dataset_name):
     
     # Configure training arguments
 
-    training_args = Seq2SeqTrainingArguments(
+    training_args = TrainingArguments(
         output_dir=f"./{FT_MODEL_NAME}",
         num_train_epochs=3, # 37,500 steps
         max_steps=1, # overwrites num_train_epochs
@@ -110,11 +111,18 @@ def fine_tune_model(base_model_name, dataset_name):
     #print("### PEFT")
     #model.print_trainable_parameters() # trainable params: 6,815,744 || all params: 8,037,076,992 || trainable%: 0.0848
     #print("###")
-        
+
+    peft_model = PeftModel.from_pretrained(
+        BASE_MODEL_NAME,
+        tokenizer=tokenizer,
+        adapter_name="lora",
+        adapter_dim=16,
+    )
+    
     # Create trainer
 
-    trainer = Seq2SeqTrainer(
-        model=model,
+    trainer = Trainer(
+        model=peft_model,
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,