|
|
''' |
|
|
Author: Qiguang Chen |
|
|
Date: 2023-01-11 10:39:26 |
|
|
LastEditors: Qiguang Chen |
|
|
LastEditTime: 2023-02-18 17:38:30 |
|
|
Description: pretrained encoder model |
|
|
|
|
|
''' |
|
|
from transformers import AutoModel, AutoConfig |
|
|
from common import utils |
|
|
|
|
|
from common.utils import InputData, HiddenData |
|
|
from model.encoder.base_encoder import BaseEncoder |
|
|
|
|
|
|
|
|
class PretrainedEncoder(BaseEncoder): |
|
|
def __init__(self, **config): |
|
|
""" init pretrained encoder |
|
|
|
|
|
Args: |
|
|
config (dict): |
|
|
encoder_name (str): pretrained model name in hugging face. |
|
|
""" |
|
|
super().__init__(**config) |
|
|
if self.config.get("_is_check_point_"): |
|
|
self.encoder = utils.instantiate(config["pretrained_model"], target="_pretrained_model_target_") |
|
|
|
|
|
else: |
|
|
self.encoder = AutoModel.from_pretrained(config["encoder_name"]) |
|
|
|
|
|
def forward(self, inputs: InputData): |
|
|
output = self.encoder(**inputs.get_inputs()) |
|
|
hidden = HiddenData(None, output.last_hidden_state) |
|
|
if self.config.get("return_with_input"): |
|
|
hidden.add_input(inputs) |
|
|
if self.config.get("return_sentence_level_hidden"): |
|
|
padding_side = self.config.get("padding_side") |
|
|
if hasattr(output, "pooler_output"): |
|
|
hidden.update_intent_hidden_state(output.pooler_output) |
|
|
elif padding_side is not None and padding_side == "left": |
|
|
hidden.update_intent_hidden_state(output.last_hidden_state[:, -1]) |
|
|
else: |
|
|
hidden.update_intent_hidden_state(output.last_hidden_state[:, 0]) |
|
|
return hidden |
|
|
|
