Merge pull request #8 from huggingface/smangrul/add-examples-fixes-docs

fixes

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,

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