From 828da005cad16bc79df6ab41c5db2ec0b287f3c1 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 7 Dec 2022 19:06:54 +0530
Subject: [PATCH] fixes

---
 src/pet/pet_model.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py
index 045bb68..e9d353c 100644
--- a/src/pet/pet_model.py
+++ b/src/pet/pet_model.py
@@ -170,10 +170,10 @@ class PETModelForSequenceClassification(PETModel):
 
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
-        self.modules_to_save = ["classifier"]
+        self.modules_to_save = ["classifier", "score"]
 
-        for name, module in self.base_model.named_children():
-            if isinstance(module, torch.nn.Linear):
+        for name, _ in self.base_model.named_children():
+            if any(module_name in name for module_name in self.modules_to_save):
                 self.cls_layer_name = name
                 break
 
@@ -237,7 +237,7 @@ class PETModelForSequenceClassification(PETModel):
             if inputs_embeds is None:
                 inputs_embeds = self.word_embeddings(input_ids)
             prompts = self.get_prompt(batch_size=batch_size)
-            inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
+            inputs_embeds = torch.cat((inputs_embeds[:, 0, :], prompts, inputs_embeds[:, 1:, :]), dim=1)
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
     def _prefix_tuning_forward(
@@ -532,10 +532,10 @@ class PETModelForTokenClassification(PETModel):
 
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
-        self.modules_to_save = ["classifier"]
+        self.modules_to_save = ["classifier", "score"]
 
-        for name, module in self.base_model.named_children():
-            if isinstance(module, torch.nn.Linear):
+        for name, _ in self.base_model.named_children():
+            if any(module_name in name for module_name in self.modules_to_save):
                 self.cls_layer_name = name
                 break
 
@@ -599,7 +599,7 @@ class PETModelForTokenClassification(PETModel):
             if inputs_embeds is None:
                 inputs_embeds = self.word_embeddings(input_ids)
             prompts = self.get_prompt(batch_size=batch_size)
-            inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
+            inputs_embeds = torch.cat((inputs_embeds[:, 0, :], prompts, inputs_embeds[:, 1:, :]), dim=1)
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
     def _prefix_tuning_forward(