From 23aecc4f6903c39fff225d9f17c504f39d9504d5 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 30 Nov 2022 14:57:23 +0530
Subject: [PATCH] fix

---
 src/pet/__init__.py        | 6 +++---
 src/pet/mapping.py         | 5 ++---
 src/pet/tuners/__init__.py | 2 +-
 src/pet/tuners/lora.py     | 7 ++++---
 4 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/pet/__init__.py b/src/pet/__init__.py
index 8843808..cce106b 100644
--- a/src/pet/__init__.py
+++ b/src/pet/__init__.py
@@ -4,9 +4,11 @@
 
 __version__ = "0.1.0.dev0"
 
-from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
 from .mapping import MODEL_TYPE_TO_PET_MODEL_MAPPING, PET_TYPE_TO_CONFIG_MAPPING, get_pet_config, get_pet_model
+from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
 from .tuners import (
+    LoRAConfig,
+    LoRAModel,
     PrefixEncoder,
     PrefixTuningConfig,
     PromptEmbedding,
@@ -15,7 +17,5 @@ from .tuners import (
     PromptEncoderReparameterizationType,
     PromptTuningConfig,
     PromptTuningInit,
-    LoRAModel,
-    LoRAConfig,
 )
 from .utils import PETConfig, PETType, PromptLearningConfig, TaskType
diff --git a/src/pet/mapping.py b/src/pet/mapping.py
index 4c82bea..4fdde10 100644
--- a/src/pet/mapping.py
+++ b/src/pet/mapping.py
@@ -1,6 +1,5 @@
 from .pet_model import PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
-from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
-from .utils import PETConfig
+from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig, LoRAConfig
 
 
 MODEL_TYPE_TO_PET_MODEL_MAPPING = {
@@ -13,7 +12,7 @@ PET_TYPE_TO_CONFIG_MAPPING = {
     "PROMPT_TUNING": PromptTuningConfig,
     "PREFIX_TUNING": PrefixTuningConfig,
     "P_TUNING": PromptEncoderConfig,
-    "LORA": PETConfig,
+    "LORA": LoRAConfig,
 }
 
 
diff --git a/src/pet/tuners/__init__.py b/src/pet/tuners/__init__.py
index d066951..22fd5c8 100644
--- a/src/pet/tuners/__init__.py
+++ b/src/pet/tuners/__init__.py
@@ -2,7 +2,7 @@
 # There's no way to ignore "F401 '...' imported but unused" warnings in this
 # module, but to preserve other warnings. So, don't check this module at all
 
+from .lora import LoRAConfig, LoRAModel
 from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType
 from .prefix_tuning import PrefixEncoder, PrefixTuningConfig
 from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit
-from .lora import LoRAModel, LoRAConfig
diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py
index 5538605..5fd7887 100644
--- a/src/pet/tuners/lora.py
+++ b/src/pet/tuners/lora.py
@@ -1,11 +1,12 @@
 # todo
-from typing import Callable, Optional
+from dataclasses import asdict, dataclass, field
+from typing import Optional
+
 import torch
 from transformers.pytorch_utils import Conv1D
-from dataclasses import dataclass, asdict, field
 
 import loralib as lora
-from loralib import mark_only_lora_as_trainable, lora_state_dict  # flake8: noqa
+from loralib import lora_state_dict, mark_only_lora_as_trainable  # noqa: F401
 
 from ..utils import PETConfig