From 2904573408caa96dbb5c46a0a7d637b3130c52e4 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Mon, 5 Dec 2022 17:11:29 +0530
Subject: [PATCH] fixes

---
 src/pet/__init__.py | 8 +++++++-
 src/pet/mapping.py  | 9 ++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/pet/__init__.py b/src/pet/__init__.py
index 734d7c2..31a3df7 100644
--- a/src/pet/__init__.py
+++ b/src/pet/__init__.py
@@ -5,7 +5,13 @@
 __version__ = "0.1.0.dev0"
 
 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 .pet_model import (
+    PETModel,
+    PETModelForCausalLM,
+    PETModelForSeq2SeqLM,
+    PETModelForSequenceClassification,
+    PETModelForTokenClassification,
+)
 from .tuners import (
     LoRAConfig,
     LoRAModel,
diff --git a/src/pet/mapping.py b/src/pet/mapping.py
index 115372e..dba6aa0 100644
--- a/src/pet/mapping.py
+++ b/src/pet/mapping.py
@@ -1,4 +1,9 @@
-from .pet_model import PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
+from .pet_model import (
+    PETModelForCausalLM,
+    PETModelForSeq2SeqLM,
+    PETModelForSequenceClassification,
+    PETModelForTokenClassification,
+)
 from .tuners import LoRAConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
 from .utils import PETType
 
@@ -7,6 +12,7 @@ MODEL_TYPE_TO_PET_MODEL_MAPPING = {
     "SEQ_CLS": PETModelForSequenceClassification,
     "SEQ_2_SEQ_LM": PETModelForSeq2SeqLM,
     "CAUSAL_LM": PETModelForCausalLM,
+    "TOKEN_CLS": PETModelForTokenClassification,
 }
 
 PET_TYPE_TO_CONFIG_MAPPING = {
@@ -32,6 +38,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
     "electra": ["query", "value"],
     "deberta-v2": ["query_proj", "value_proj"],
     "deberta": ["in_proj"],
+    "layoutlm": ["query", "value"],
 }