From bba5b36d332b2b3c0297c85cd7d07e5a50e3f86c Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Tue, 3 Sep 2024 19:40:53 -0700
Subject: [PATCH] Don't import PEFT unless necessary

---
 modules/LoRA.py     |  4 +++-
 modules/training.py | 18 ++++++++++--------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/modules/LoRA.py b/modules/LoRA.py
index 117022cf..4fd144ba 100644
--- a/modules/LoRA.py
+++ b/modules/LoRA.py
@@ -1,7 +1,6 @@
 from pathlib import Path
 
 import torch
-from peft import PeftModel
 from transformers import is_torch_xpu_available
 
 import modules.shared as shared
@@ -85,6 +84,9 @@ def add_lora_autogptq(lora_names):
 
 
 def add_lora_transformers(lora_names):
+
+    from peft import PeftModel
+
     prior_set = set(shared.lora_names)
     added_set = set(lora_names) - prior_set
     removed_set = prior_set - set(lora_names)
diff --git a/modules/training.py b/modules/training.py
index b003fc8c..11c4b8c5 100644
--- a/modules/training.py
+++ b/modules/training.py
@@ -18,14 +18,6 @@ import gradio as gr
 import torch
 import transformers
 from datasets import Dataset, load_dataset
-from peft import (
-    LoraConfig,
-    get_peft_model,
-    prepare_model_for_kbit_training,
-    set_peft_model_state_dict
-)
-from peft.utils.other import \
-    TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING as model_to_lora_modules
 from transformers import is_torch_xpu_available
 from transformers.models.auto.modeling_auto import (
     MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
@@ -292,6 +284,16 @@ def calc_trainable_parameters(model):
 
 def do_train(lora_name: str, always_override: bool, q_proj_en: bool, v_proj_en: bool, k_proj_en: bool, o_proj_en: bool, gate_proj_en: bool, down_proj_en: bool, up_proj_en: bool, save_steps: int, micro_batch_size: int, batch_size: int, epochs: int, learning_rate: str, lr_scheduler_type: str, lora_rank: int, lora_alpha: int, lora_dropout: float, cutoff_len: int, dataset: str, eval_dataset: str, format: str, eval_steps: int, raw_text_file: str, overlap_len: int, newline_favor_len: int, higher_rank_limit: bool, warmup_steps: int, optimizer: str, hard_cut_string: str, train_only_after: str, stop_at_loss: float, add_eos_token: bool, min_chars: int, report_to: str):
 
+    from peft import (
+        LoraConfig,
+        get_peft_model,
+        prepare_model_for_kbit_training,
+        set_peft_model_state_dict
+    )
+    from peft.utils.other import \
+        TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING as \
+        model_to_lora_modules
+
     global WANT_INTERRUPT
     WANT_INTERRUPT = False