Spaces:

maratale
/

NLP-BERT

Sleeping

NLP-BERT / src /LSTMattention_classes.py

d25442c about 1 year ago

3.84 kB

import sys
import os
import torch
from torch.utils.data import DataLoader, TensorDataset
import torch.nn.functional as F
import torch.nn as nn
import torchutils as tu
from dataclasses import dataclass
from typing import Union
import numpy as np


@dataclass
class Config:
    n_layers: int
    embedding_size: int
    hidden_size: int
    vocab_size: int
    device: str
    seq_len: int
    bidirectional: Union[bool, int] = False


class BahdanauAttention(nn.Module):
    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.linear_key = nn.Linear(hidden_size, hidden_size)
        self.linear_query = nn.Linear(hidden_size, hidden_size)
        self.cls = nn.Linear(hidden_size, 1)
        self.tanh = nn.Tanh()

    def forward(self, lstm_outputs, final_hidden):
        # print(f"LSTM output shape: {lstm_outputs.shape}")
        # print(f"Final_hidden shape: {final_hidden.shape}")
        keys = self.linear_key(lstm_outputs)  # (batch_size, seq_len, hidden_size)
        # print(f"After linear keys shape: {keys.shape}")
        query = self.linear_query(final_hidden)  # (batch_size, hidden_size)
        query = query.unsqueeze(1)  # (batch_size, 1, hidden_size)
        # print(f"After linear query shape: {query.shape}")
        x = self.tanh(keys + query)  # (batch_size, seq_len, hidden_size)
        # print(f"After + X shape: {x.shape}")
        x = self.cls(x)  # (batch_size, seq_len, 1)
        # print(f"After cls x shape: {x.shape}")
        x = x.squeeze(-1)  # (batch_size, seq_len)
        # print(f"After squeeze x shape: {x.shape}")
        attention_weights = F.softmax(x, dim=-1)  # (batch_size, seq_len)
        # print(f"Attention weights shape: {attention_weights.shape}")
        attention_weights_bmm = attention_weights.unsqueeze(
            1
        )  # (batch_size, 1, seq_len)
        # print(f"Attention weights for bmm shape: {attention_weights_bmm.shape}")

        # bmm : (batch_size, 1, seq_len) * (batch_size, seq_len, hidden_size) = (batch_size, 1, hidden_size)
        context = torch.bmm(attention_weights_bmm, keys)  # (batch_size, 1, hidden_size)
        # print(f"Context shape: {context.shape}")
        context = context.squeeze(1)
        # print(f"Context final shape: {context.shape}")

        return context, attention_weights


class LSTMBahdanauAttention(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()

        # инициализируем конфиг
        self.config = config
        self.seq_len = self.config.seq_len
        self.vocab_size = self.config.vocab_size
        self.hidden_size = self.config.hidden_size
        self.emb_size = self.config.embedding_size
        self.n_layers = self.config.n_layers
        self.device = self.config.device
        self.bidirectional = bool(self.config.bidirectional)

        self.embedding = nn.Embedding.from_pretrained(
            torch.FloatTensor(np.zeros((self.vocab_size, self.emb_size)))
        )
        self.lstm = nn.LSTM(self.emb_size, self.hidden_size, batch_first=True)
        self.bidirect_factor = 2 if self.bidirectional == 1 else 1
        self.attn = BahdanauAttention(self.hidden_size)
        self.clf = nn.Sequential(
            nn.Linear(self.hidden_size, 128), nn.Dropout(), nn.Tanh(), nn.Linear(128, 1)
        )

    def model_description(self):
        direction = "bidirect" if self.bidirectional else "onedirect"
        return f"rnn_{direction}_{self.n_layers}"

    def forward(self, x):
        embeddings = self.embedding(x)
        outputs, (h_n, _) = self.lstm(embeddings)
        # att_hidden, att_weights = self.attn(outputs, h_n[-1].squeeze(0))
        att_hidden, att_weights = self.attn(outputs, h_n[-1])
        out = self.clf(att_hidden)
        return out, att_weights

	import sys
	import os
	import torch
	from torch.utils.data import DataLoader, TensorDataset
	import torch.nn.functional as F
	import torch.nn as nn
	import torchutils as tu
	from dataclasses import dataclass
	from typing import Union
	import numpy as np


	@dataclass
	class Config:
	n_layers: int
	embedding_size: int
	hidden_size: int
	vocab_size: int
	device: str
	seq_len: int
	bidirectional: Union[bool, int] = False


	class BahdanauAttention(nn.Module):
	def __init__(self, hidden_size: int) -> None:
	super().__init__()
	self.hidden_size = hidden_size
	self.linear_key = nn.Linear(hidden_size, hidden_size)
	self.linear_query = nn.Linear(hidden_size, hidden_size)
	self.cls = nn.Linear(hidden_size, 1)
	self.tanh = nn.Tanh()

	def forward(self, lstm_outputs, final_hidden):
	# print(f"LSTM output shape: {lstm_outputs.shape}")
	# print(f"Final_hidden shape: {final_hidden.shape}")
	keys = self.linear_key(lstm_outputs) # (batch_size, seq_len, hidden_size)
	# print(f"After linear keys shape: {keys.shape}")
	query = self.linear_query(final_hidden) # (batch_size, hidden_size)
	query = query.unsqueeze(1) # (batch_size, 1, hidden_size)
	# print(f"After linear query shape: {query.shape}")
	x = self.tanh(keys + query) # (batch_size, seq_len, hidden_size)
	# print(f"After + X shape: {x.shape}")
	x = self.cls(x) # (batch_size, seq_len, 1)
	# print(f"After cls x shape: {x.shape}")
	x = x.squeeze(-1) # (batch_size, seq_len)
	# print(f"After squeeze x shape: {x.shape}")
	attention_weights = F.softmax(x, dim=-1) # (batch_size, seq_len)
	# print(f"Attention weights shape: {attention_weights.shape}")
	attention_weights_bmm = attention_weights.unsqueeze(
	1
	) # (batch_size, 1, seq_len)
	# print(f"Attention weights for bmm shape: {attention_weights_bmm.shape}")

	# bmm : (batch_size, 1, seq_len) * (batch_size, seq_len, hidden_size) = (batch_size, 1, hidden_size)
	context = torch.bmm(attention_weights_bmm, keys) # (batch_size, 1, hidden_size)
	# print(f"Context shape: {context.shape}")
	context = context.squeeze(1)
	# print(f"Context final shape: {context.shape}")

	return context, attention_weights


	class LSTMBahdanauAttention(nn.Module):
	def __init__(self, config) -> None:
	super().__init__()

	# инициализируем конфиг
	self.config = config
	self.seq_len = self.config.seq_len
	self.vocab_size = self.config.vocab_size
	self.hidden_size = self.config.hidden_size
	self.emb_size = self.config.embedding_size
	self.n_layers = self.config.n_layers
	self.device = self.config.device
	self.bidirectional = bool(self.config.bidirectional)

	self.embedding = nn.Embedding.from_pretrained(
	torch.FloatTensor(np.zeros((self.vocab_size, self.emb_size)))
	)
	self.lstm = nn.LSTM(self.emb_size, self.hidden_size, batch_first=True)
	self.bidirect_factor = 2 if self.bidirectional == 1 else 1
	self.attn = BahdanauAttention(self.hidden_size)
	self.clf = nn.Sequential(
	nn.Linear(self.hidden_size, 128), nn.Dropout(), nn.Tanh(), nn.Linear(128, 1)
	)

	def model_description(self):
	direction = "bidirect" if self.bidirectional else "onedirect"
	return f"rnn_{direction}_{self.n_layers}"

	def forward(self, x):
	embeddings = self.embedding(x)
	outputs, (h_n, _) = self.lstm(embeddings)
	# att_hidden, att_weights = self.attn(outputs, h_n[-1].squeeze(0))
	att_hidden, att_weights = self.attn(outputs, h_n[-1])
	out = self.clf(att_hidden)
	return out, att_weights