Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from transformers import AutoModelForSeq2SeqLM, AutoModelForMaskedLM | |
| from torch.nn import TransformerDecoder, TransformerDecoderLayer | |
| import math | |
| class MultiHeadAttention(nn.Module): | |
| def __init__(self, n_heads, d_model, dropout=0.1): | |
| super(MultiHeadAttention, self).__init__() | |
| self.n_heads = n_heads | |
| self.d_model = d_model | |
| self.dropout = nn.Dropout(p=dropout) | |
| self.h = nn.Parameter(torch.empty(n_heads, d_model // n_heads), requires_grad=True) | |
| self.linear_q = nn.Linear(d_model, d_model) | |
| self.linear_k = nn.Linear(d_model, d_model) | |
| self.linear_v = nn.Linear(d_model, d_model) | |
| self.norm = nn.LayerNorm(d_model) | |
| self.fc = nn.Linear(d_model, d_model) | |
| def forward(self, query, key, value): | |
| # 计算query、key和value的线性变换 | |
| q = self.linear_q(query) | |
| k = self.linear_k(key) | |
| v = self.linear_v(value) | |
| # 分割成多个头 | |
| q, k, v = q.chunk(self.n_heads, dim=-1), k.chunk(self.n_heads, dim=-1), v.chunk(self.n_heads, dim=-1) | |
| # 计算注意力分数 | |
| scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.h.size(1)) | |
| # 应用softmax函数 | |
| weights = F.softmax(scores, dim=-1) | |
| # 加权求和 | |
| outputs = torch.matmul(weights, v) | |
| # 反卷积操作 | |
| outputs = outputs.transpose(1, 2).contiguous() | |
| # 再次应用线性变换 | |
| outputs = self.fc(outputs) | |
| # 应用层归一化 | |
| return self.norm(outputs) | |
| class SimplifiedTransformerDecoderLayer(nn.Module): | |
| def __init__(self, d_model, output_dim, vocab_size, nhead, dim_feedforward, dropout=0.1): | |
| super(SimplifiedTransformerDecoderLayer, self).__init__() | |
| # self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) | |
| self.self_attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=8, kdim=d_model, vdim=d_model, | |
| batch_first=True) | |
| self.linear1 = nn.Linear(d_model, dim_feedforward) | |
| self.dropout = nn.Dropout(dropout) | |
| self.linear2 = nn.Linear(dim_feedforward, d_model) | |
| self.linear_vocab = nn.Linear(d_model, 1) | |
| self.norm1 = nn.LayerNorm(d_model) | |
| self.relu = nn.ReLU() | |
| def forward(self, tgt, tar): | |
| tar1 = self.self_attn(query=tar, key=tar, value=tar)[0] | |
| tar = tar + tar1 | |
| tar = self.norm1(tar) | |
| tgt2 = self.self_attn(query=tar, key=tgt, value=tgt)[0] | |
| tgt3 = tar + self.dropout(tgt2) | |
| tgt3 = self.norm1(tgt3) | |
| # 线性层 + ReLU + Dropout | |
| tgt2 = self.linear2(self.dropout(self.relu(self.linear1(tgt3)))) | |
| tar = tar + self.dropout(tgt2) | |
| return tar | |
| class BirdModel_Attention_lstm(nn.Module): | |
| def __init__(self, model, key): | |
| super(BirdModel_Attention_lstm, self).__init__() | |
| self.model = model | |
| self.decoder = SimplifiedTransformerDecoderLayer(d_model=768, output_dim=1, vocab_size=1, nhead=8, | |
| dim_feedforward=1024) | |
| self.linear_vocab = nn.Linear(1024, 2) | |
| self.relu = nn.ReLU() | |
| self.lstm = nn.LSTM(768, 128, 5, | |
| bidirectional=True, batch_first=True, dropout=0.1) | |
| self.maxpool = nn.MaxPool1d(32) | |
| self.fc = nn.Linear(128 * 2 + 768, 2) | |
| if key: | |
| print("begin-------Training from scratch !!!") | |
| for param in self.model.parameters(): | |
| param.requires_grad = False | |
| else: | |
| print("begin-------fine train !!!") | |
| for name, param in reversed(list(model.named_parameters())): | |
| if "cls" in name: | |
| param.requires_grad = False | |
| def forward(self, input_ids, masks): | |
| outputs = self.model.bert(input_ids, attention_mask=masks).last_hidden_state | |
| tar = torch.zeros(outputs.size(0), 256, 768) | |
| tar = self.decoder(outputs, tar) | |
| tar = self.decoder(outputs, tar) | |
| tar = self.decoder(outputs, tar) | |
| tar = self.decoder(outputs, tar) | |
| out, _ = self.lstm(tar) | |
| out = torch.cat((tar, out), 2) | |
| out = self.relu(out) | |
| out = out.permute(0, 2, 1) | |
| pool = nn.MaxPool1d(out.size(-1)) | |
| out = pool(out).squeeze() | |
| out = self.fc(out) | |
| return out | |