Spaces:
Sleeping
Sleeping
| import json | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| # Hyperparameters used to train pico_model.pth, they must match the checkpoint | |
| D_MODEL = 512 | |
| NHEAD = 8 | |
| DIM_FF = 4 * D_MODEL | |
| NUM_ENCODER_LAYERS = 3 | |
| NUM_DECODER_LAYERS = 3 | |
| MAX_LEN = 122 # 120 characters + start and end symbols | |
| class Embedding(nn.Module): | |
| def __init__(self, vocab_size, d_model, dropout=0.1, max_len=MAX_LEN): | |
| super().__init__() | |
| self.d_model = d_model | |
| self.input_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0) | |
| # Fixed sinusoidal positional encoding, not learned | |
| self.pos_embedding = nn.Embedding.from_pretrained( | |
| self.pos_encoding(max_len, d_model), freeze=True) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, X): | |
| pos_mat = torch.arange(X.size(dim=-1), device=X.device) | |
| X = self.input_embedding(X) * math.sqrt(self.d_model) | |
| X += self.pos_embedding(pos_mat) | |
| return self.dropout(X) | |
| def pos_encoding(self, max_len, d_model): | |
| dividend = torch.arange(max_len).unsqueeze(0).T | |
| divisor = torch.pow(10000.0, torch.arange(0, d_model, 2) / d_model) | |
| angles = dividend / divisor | |
| pe = torch.zeros((max_len, d_model)) | |
| pe[:, 0::2] = torch.sin(angles) | |
| pe[:, 1::2] = torch.cos(angles) | |
| return pe | |
| class Translator(nn.Module): | |
| def __init__(self, vocab_size, d_model=D_MODEL, nhead=NHEAD, | |
| num_encoder_layers=NUM_ENCODER_LAYERS, | |
| num_decoder_layers=NUM_DECODER_LAYERS, | |
| dim_feedforward=DIM_FF, dropout=0.1, max_len=MAX_LEN): | |
| super().__init__() | |
| self.embedding = Embedding(vocab_size, d_model, max_len=max_len) | |
| self.transformer = nn.Transformer( | |
| d_model, nhead, num_encoder_layers, num_decoder_layers, | |
| dim_feedforward, dropout, batch_first=True) | |
| self.fc = nn.Linear(d_model, vocab_size) | |
| self.fc.weight = self.embedding.input_embedding.weight | |
| def load(model_path, vocab_path): | |
| with open(vocab_path, encoding='utf-8') as f: | |
| token2idx = json.load(f) | |
| idx2token = {idx: tok for tok, idx in token2idx.items()} | |
| model = Translator(vocab_size=len(token2idx)) | |
| model.load_state_dict(torch.load(model_path, map_location='cpu')) | |
| model.eval() | |
| return model, token2idx, idx2token | |
| def translate(model, token2idx, idx2token, sentence): | |
| """Greedy decoding, one character at a time""" | |
| chars = list(sentence.strip())[:MAX_LEN - 2] | |
| if not chars: | |
| return '' | |
| src = torch.LongTensor( | |
| [[token2idx['<s>']] + [token2idx.get(c, token2idx['<unk>']) | |
| for c in chars] + [token2idx['</s>']]]) | |
| memory = model.transformer.encoder(model.embedding(src)) | |
| tgt = torch.LongTensor([[token2idx['<s>']]]) | |
| max_len = min(len(chars) + 20, MAX_LEN) | |
| for _ in range(max_len - 1): | |
| tgt_mask = nn.Transformer.generate_square_subsequent_mask( | |
| tgt.size(dim=1)) | |
| out = model.transformer.decoder( | |
| model.embedding(tgt), memory, tgt_mask=tgt_mask) | |
| next_char = torch.argmax(model.fc(out[:, -1, :]), dim=-1).unsqueeze(-1) | |
| tgt = torch.cat((tgt, next_char), dim=1) | |
| if next_char.item() == token2idx['</s>']: | |
| break | |
| tokens = [idx2token[i] for i in tgt[0, 1:].tolist()] | |
| if tokens and tokens[-1] == '</s>': | |
| tokens = tokens[:-1] | |
| return ''.join(tokens) | |