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
+from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_int8_training
 from transformers import AutoModelForCausalLM, AutoTokenizer, Seq2SeqTrainer, Seq2SeqTrainingArguments, pipeline
 
 ACTION_1 = "Prompt base model"
@@ -46,12 +46,14 @@ def fine_tune_model(base_model_name, dataset_name):
     
     model, tokenizer = load_model(base_model_name)
 
+    model.print_trainable_parameters()
+
     print("### Model")
     print(model)
     print("### Tokenizer")
     print(tokenizer)
     print("###")
-    
+        
     # Pre-process dataset
     
     def preprocess(examples):
@@ -91,14 +93,22 @@ def fine_tune_model(base_model_name, dataset_name):
     print(training_args)
     print("###")
 
-    # PEFT
+    # PEFT https://www.philschmid.de/fine-tune-flan-t5-peft
 
-    #peft_config = LoraConfig(
-    #    r=8,
-    #    bias="none",
-    #    task_type="CAUSAL_LM",
-    #)
+    lora_config = LoraConfig(
+        r=16,
+        # TODO
+        #bias="none",
+        #lora_alpha=32,
+        #lora_dropout=0.05,
+        #target_modules=["q", "v"],
+        task_type=TaskType.SEQ_2_SEQ_LM,
+    )
     
+    model = prepare_model_for_int8_training(model)
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+        
     # Create trainer
 
     trainer = Seq2SeqTrainer(
@@ -106,17 +116,16 @@ 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
     )
 
     # Train model
     
-    trainer.train()
+    #trainer.train()
 
     # Push tokenizer to HF
 
-    tokenizer.push_to_hub(FT_MODEL_NAME)
+    #tokenizer.push_to_hub(FT_MODEL_NAME)
     
 def prompt_model(model_name, system_prompt, user_prompt, sql_context):
     pipe = pipeline("text-generation",