models.py

import torch
import torch.nn as nn
from transformers import (
    AutoConfig,
    AutoModel,
    PreTrainedModel,
    T5ForConditionalGeneration,
)


class ReactionT5Yield(nn.Module):
    def __init__(self, cfg, config_path=None, pretrained=False):
        super().__init__()
        self.cfg = cfg
        if config_path is None:
            self.config = AutoConfig.from_pretrained(
                self.cfg.pretrained_model_name_or_path, output_hidden_states=True
            )
        else:
            self.config = torch.load(config_path, weights_only=False)
        if pretrained:
            self.model = AutoModel.from_pretrained(
                self.cfg.pretrained_model_name_or_path
            )
        else:
            self.model = AutoModel.from_config(self.config)
        self.model.resize_token_embeddings(len(self.cfg.tokenizer))
        self.fc_dropout1 = nn.Dropout(self.cfg.fc_dropout)
        self.fc1 = nn.Linear(self.config.hidden_size, self.config.hidden_size // 2)
        self.fc_dropout2 = nn.Dropout(self.cfg.fc_dropout)

        self.fc2 = nn.Linear(self.config.hidden_size, self.config.hidden_size // 2)
        self.fc3 = nn.Linear(self.config.hidden_size // 2 * 2, self.config.hidden_size)
        self.fc4 = nn.Linear(self.config.hidden_size, self.config.hidden_size)
        self.fc5 = nn.Linear(self.config.hidden_size, 1)

        self._init_weights(self.fc1)
        self._init_weights(self.fc2)
        self._init_weights(self.fc3)
        self._init_weights(self.fc4)
        self._init_weights(self.fc5)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.01)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.01)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def forward(self, inputs):
        encoder_outputs = self.model.encoder(**inputs)
        encoder_hidden_states = encoder_outputs[0]
        outputs = self.model.decoder(
            input_ids=torch.full(
                (inputs["input_ids"].size(0), 1),
                self.config.decoder_start_token_id,
                dtype=torch.long,
                device=inputs["input_ids"].device,
            ),
            encoder_hidden_states=encoder_hidden_states,
        )
        last_hidden_states = outputs[0]
        output1 = self.fc1(
            self.fc_dropout1(last_hidden_states).view(-1, self.config.hidden_size)
        )
        output2 = self.fc2(
            encoder_hidden_states[:, 0, :].view(-1, self.config.hidden_size)
        )
        output = self.fc3(self.fc_dropout2(torch.hstack((output1, output2))))
        output = self.fc4(output)
        output = self.fc5(output)
        return output

    def generate_embedding(self, inputs):
        encoder_outputs = self.model.encoder(**inputs)
        encoder_hidden_states = encoder_outputs[0]
        outputs = self.model.decoder(
            input_ids=torch.full(
                (inputs["input_ids"].size(0), 1),
                self.config.decoder_start_token_id,
                dtype=torch.long,
                device=inputs["input_ids"].device,
            ),
            encoder_hidden_states=encoder_hidden_states,
        )
        last_hidden_states = outputs[0]
        output1 = self.fc1(
            self.fc_dropout1(last_hidden_states).view(-1, self.config.hidden_size)
        )
        output2 = self.fc2(
            encoder_hidden_states[:, 0, :].view(-1, self.config.hidden_size)
        )
        return torch.hstack((output1, output2))


class ReactionT5Yield2(PreTrainedModel):
    config_class = AutoConfig

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.model = T5ForConditionalGeneration.from_pretrained(
            self.config._name_or_path
        )
        self.model.resize_token_embeddings(self.config.vocab_size)
        self.fc1 = nn.Linear(self.config.hidden_size, self.config.hidden_size // 2)
        self.fc2 = nn.Linear(self.config.hidden_size, self.config.hidden_size // 2)
        self.fc3 = nn.Linear(self.config.hidden_size // 2 * 2, self.config.hidden_size)
        self.fc4 = nn.Linear(self.config.hidden_size, self.config.hidden_size)
        self.fc5 = nn.Linear(self.config.hidden_size, 1)

        self._init_weights(self.fc1)
        self._init_weights(self.fc2)
        self._init_weights(self.fc3)
        self._init_weights(self.fc4)
        self._init_weights(self.fc5)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.01)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.01)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def forward(self, inputs):
        encoder_outputs = self.model.encoder(**inputs)
        encoder_hidden_states = encoder_outputs[0]
        outputs = self.model.decoder(
            input_ids=torch.full(
                (inputs["input_ids"].size(0), 1),
                self.config.decoder_start_token_id,
                dtype=torch.long,
                device=inputs["input_ids"].device,
            ),
            encoder_hidden_states=encoder_hidden_states,
        )
        last_hidden_states = outputs[0]
        output1 = self.fc1(last_hidden_states.view(-1, self.config.hidden_size))
        output2 = self.fc2(
            encoder_hidden_states[:, 0, :].view(-1, self.config.hidden_size)
        )
        output = self.fc3(torch.hstack((output1, output2)))
        output = self.fc4(output)
        output = self.fc5(output)
        return output * 100

    def generate_embedding(self, inputs):
        encoder_outputs = self.model.encoder(**inputs)
        encoder_hidden_states = encoder_outputs[0]
        outputs = self.model.decoder(
            input_ids=torch.full(
                (inputs["input_ids"].size(0), 1),
                self.config.decoder_start_token_id,
                dtype=torch.long,
                device=inputs["input_ids"].device,
            ),
            encoder_hidden_states=encoder_hidden_states,
        )
        last_hidden_states = outputs[0]
        output1 = self.fc1(last_hidden_states.view(-1, self.config.hidden_size))
        output2 = self.fc2(
            encoder_hidden_states[:, 0, :].view(-1, self.config.hidden_size)
        )
        return torch.hstack((output1, output2))