Update app.py

app.py

@@ -4,8 +4,8 @@ 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, prepare_model_for_kbit_training
-from transformers import AutoModelForCausalLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, Trainer, TrainingArguments, pipeline
+#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"
 ACTION_2 = "Fine-tune base model"
@@ -79,7 +79,7 @@ def fine_tune_model(base_model_name, dataset_name):
     
     # Configure training arguments
 
-    training_args = TrainingArguments(
+    training_args = Seq2SeqTrainingArguments(
         output_dir=f"./{FT_MODEL_NAME}",
         num_train_epochs=3, # 37,500 steps
         max_steps=1, # overwrites num_train_epochs
@@ -93,27 +93,27 @@ def fine_tune_model(base_model_name, dataset_name):
 
     # PEFT
 
-    lora_config = LoraConfig(
-        r=16,
+    #lora_config = LoraConfig(
+    #    r=16,
         # TODO https://www.philschmid.de/fine-tune-flan-t5-peft
         #bias="none",
         #lora_alpha=32,
         #lora_dropout=0.05,
         #target_modules=["q", "v"],
-        task_type=TaskType.SEQ_2_SEQ_LM,
-    )
+    #    task_type=TaskType.SEQ_2_SEQ_LM,
+    #)
 
-    model = prepare_model_for_kbit_training(model)
+    #model = prepare_model_for_kbit_training(model)
     
-    model = get_peft_model(model, lora_config)
+    #model = get_peft_model(model, lora_config)
     
-    print("### PEFT")
-    model.print_trainable_parameters() # trainable params: 6,815,744 || all params: 8,037,076,992 || trainable%: 0.0848
-    print("###")
+    #print("### PEFT")
+    #model.print_trainable_parameters() # trainable params: 6,815,744 || all params: 8,037,076,992 || trainable%: 0.0848
+    #print("###")
         
     # Create trainer
 
-    trainer = Trainer(
+    trainer = Seq2SeqTrainer(
         model=model,
         args=training_args,
         train_dataset=train_dataset,
@@ -152,7 +152,7 @@ def prompt_model(model_name, system_prompt, user_prompt, sql_context):
     return result
 
 def load_model(model_name):
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     tokenizer.pad_token = tokenizer.eos_token