fix

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

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,
 }
 
 

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

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