This commit is contained in:
Sourab Mangrulkar
2022-12-26 17:19:26 +05:30
parent 867b66bc75
commit c0c195a37c
2 changed files with 3 additions and 3 deletions
+2 -2
View File
@@ -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((inputs_embeds[:, 0, :], prompts, inputs_embeds[:, 1:, :]), dim=1)
inputs_embeds = torch.cat((inputs_embeds[:, 0, :].unsqueeze(1), prompts, inputs_embeds[:, 1:, :]), dim=1)
return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
def _prefix_tuning_forward(
@@ -680,7 +680,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((inputs_embeds[:, 0, :], prompts, inputs_embeds[:, 1:, :]), dim=1)
inputs_embeds = torch.cat((inputs_embeds[:, 0, :].unsqueeze(1), prompts, inputs_embeds[:, 1:, :]), dim=1)
return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
def _prefix_tuning_forward(
+1 -1
View File
@@ -1,7 +1,6 @@
from loralib import lora_state_dict
from .config import PETType
from ..mapping import get_pet_model
def get_pet_model_state_dict(model, state_dict=None):
@@ -61,6 +60,7 @@ def pet_model_load_and_dispatch(model, pet_model_state_dict, pet_config, max_mem
"""
from accelerate import infer_auto_device_map, dispatch_model
from accelerate.hooks import remove_hook_from_submodules, AlignDevicesHook, add_hook_to_module
from ..mapping import get_pet_model
remove_hook_from_submodules(model)
model = get_pet_model(model, pet_config)