Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import torch as T | |
| from tqdm import tqdm | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from components.k_lstm import K_LSTM | |
| from components.attention import Attention | |
| from data_utils import DatasetUtils | |
| from diac_utils import flat2_3head, flat_2_3head | |
| class DiacritizerD2(nn.Module): | |
| def __init__(self, config): | |
| super(DiacritizerD2, self).__init__() | |
| self.max_word_len = config["train"]["max-word-len"] | |
| self.max_sent_len = config["train"]["max-sent-len"] | |
| self.char_embed_dim = config["train"]["char-embed-dim"] | |
| self.final_dropout_p = config["train"]["final-dropout"] | |
| self.sent_dropout_p = config["train"]["sent-dropout"] | |
| self.diac_dropout_p = config["train"]["diac-dropout"] | |
| self.vertical_dropout = config['train']['vertical-dropout'] | |
| self.recurrent_dropout = config['train']['recurrent-dropout'] | |
| self.recurrent_dropout_mode = config['train'].get('recurrent-dropout-mode', 'gal_tied') | |
| self.recurrent_activation = config['train'].get('recurrent-activation', 'sigmoid') | |
| self.sent_lstm_units = config["train"]["sent-lstm-units"] | |
| self.word_lstm_units = config["train"]["word-lstm-units"] | |
| self.decoder_units = config["train"]["decoder-units"] | |
| self.sent_lstm_layers = config["train"]["sent-lstm-layers"] | |
| self.word_lstm_layers = config["train"]["word-lstm-layers"] | |
| self.cell = config['train'].get('rnn-cell', 'lstm') | |
| self.num_layers = config["train"].get("num-layers", 2) | |
| self.RNN_Layer = K_LSTM | |
| self.batch_first = config['train'].get('batch-first', True) | |
| self.device = 'cuda' if T.cuda.is_available() else 'cpu' | |
| self.num_classes = 15 | |
| def build(self, wembs: T.Tensor, abjad_size: int): | |
| self.closs = F.cross_entropy | |
| self.bloss = F.binary_cross_entropy_with_logits | |
| rnn_kargs = dict( | |
| recurrent_dropout_mode=self.recurrent_dropout_mode, | |
| recurrent_activation=self.recurrent_activation, | |
| ) | |
| self.sent_lstm = self.RNN_Layer( | |
| input_size=300, | |
| hidden_size=self.sent_lstm_units, | |
| num_layers=self.sent_lstm_layers, | |
| bidirectional=True, | |
| vertical_dropout=self.vertical_dropout, | |
| recurrent_dropout=self.recurrent_dropout, | |
| batch_first=self.batch_first, | |
| **rnn_kargs, | |
| ) | |
| self.word_lstm = self.RNN_Layer( | |
| input_size=self.sent_lstm_units * 2 + self.char_embed_dim, | |
| hidden_size=self.word_lstm_units, | |
| num_layers=self.word_lstm_layers, | |
| bidirectional=True, | |
| vertical_dropout=self.vertical_dropout, | |
| recurrent_dropout=self.recurrent_dropout, | |
| batch_first=self.batch_first, | |
| return_states=True, | |
| **rnn_kargs, | |
| ) | |
| self.char_embs = nn.Embedding( | |
| abjad_size, | |
| self.char_embed_dim, | |
| padding_idx=0, | |
| ) | |
| self.attention = Attention( | |
| kind="dot", | |
| query_dim=self.word_lstm_units * 2, | |
| input_dim=self.sent_lstm_units * 2, | |
| ) | |
| self.word_embs = T.tensor(wembs).clone().to(dtype=T.float32) | |
| self.word_embs = self.word_embs.to(self.device) | |
| self.classifier = nn.Linear(self.attention.Dout + self.word_lstm_units * 2, self.num_classes) | |
| self.dropout = nn.Dropout(self.final_dropout_p) | |
| def forward(self, sents, words, labels=None, subword_lengths=None): | |
| #^ sents : [b ts] | |
| #^ words : [b ts tw] | |
| #^ labels: [b ts tw] | |
| max_words = min(self.max_sent_len, sents.shape[1]) | |
| word_mask = words.ne(0.).float() | |
| #^ word_mask: [b ts tw] | |
| if self.training: | |
| q = 1.0 - self.sent_dropout_p | |
| sdo = T.bernoulli(T.full(sents.shape, q)) | |
| sents_do = sents * sdo.long() | |
| #^ sents_do : [b ts] ; DO(ts) | |
| wembs = self.word_embs[sents_do] | |
| #^ wembs : [b ts dw] ; DO(ts) | |
| else: | |
| wembs = self.word_embs[sents] | |
| #^ wembs : [b ts dw] | |
| sent_enc = self.sent_lstm(wembs.to(self.device)) | |
| #^ sent_enc : [b ts dwe] | |
| sentword_do = sent_enc.unsqueeze(2) | |
| #^ sentword_do : [b ts _ dwe] | |
| sentword_do = self.dropout(sentword_do * word_mask.unsqueeze(-1)) | |
| #^ sentword_do : [b ts tw dwe] | |
| word_index = words.view(-1, self.max_word_len) | |
| #^ word_index: [b*ts tw]? | |
| cembs = self.char_embs(word_index) | |
| #^ cembs : [b*ts tw dc] | |
| sentword_do = sentword_do.view(-1, self.max_word_len, self.sent_lstm_units * 2) | |
| #^ sentword_do : [b*ts tw dwe] | |
| char_embs = T.cat([cembs, sentword_do], dim=-1) | |
| #^ char_embs : [b*ts tw dcw] ; dcw = dc + dwe | |
| char_enc, _ = self.word_lstm(char_embs) | |
| #^ char_enc: [b*ts tw dce] | |
| char_enc_reshaped = char_enc.view(-1, max_words, self.max_word_len, self.word_lstm_units * 2) | |
| # #^ char_enc: [b ts tw dce] | |
| omit_self_mask = (1.0 - T.eye(max_words)).unsqueeze(0).to(self.device) | |
| attn_enc, attn_map = self.attention(char_enc_reshaped, sent_enc, word_mask.bool(), prejudice_mask=omit_self_mask) | |
| # # #^ attn_enc: [b ts tw dae] | |
| attn_enc = attn_enc.reshape(-1, self.max_word_len, self.attention.Dout) | |
| # #^ attn_enc: [b*ts tw dae] | |
| final_vec = T.cat([attn_enc, char_enc], dim=-1) | |
| diac_out = self.classifier(self.dropout(final_vec)) | |
| #^ diac_out: [b*ts tw 7] | |
| diac_out = diac_out.view(-1, max_words, self.max_word_len, self.num_classes) | |
| #^ diac_out: [b ts tw 7] | |
| if not self.batch_first: | |
| diac_out = diac_out.swapaxes(1, 0) | |
| return diac_out | |
| def step(self, xt, yt, mask=None): | |
| xt[1] = xt[1].to(self.device) | |
| xt[2] = xt[2].to(self.device) | |
| yt = yt.to(self.device) | |
| #^ yt: [b ts tw] | |
| diac, _ = self(*xt) | |
| loss = self.closs(diac.view(-1, self.num_classes), yt.view(-1)) | |
| return loss | |
| def predict(self, dataloader): | |
| training = self.training | |
| self.eval() | |
| preds = {'haraka': [], 'shadda': [], 'tanween': []} | |
| print("> Predicting...") | |
| for inputs, _ in tqdm(dataloader, total=len(dataloader)): | |
| inputs[0] = inputs[0].to(self.device) | |
| inputs[1] = inputs[1].to(self.device) | |
| diac = self(*inputs) | |
| output = np.argmax(T.softmax(diac.detach(), dim=-1).cpu().numpy(), axis=-1) | |
| #^ [b ts tw] | |
| haraka, tanween, shadda = flat_2_3head(output) | |
| preds['haraka'].extend(haraka) | |
| preds['tanween'].extend(tanween) | |
| preds['shadda'].extend(shadda) | |
| self.train(training) | |
| return ( | |
| np.array(preds['haraka']), | |
| np.array(preds["tanween"]), | |
| np.array(preds["shadda"]), | |
| ) | |
| class DiacritizerD3(nn.Module): | |
| def __init__(self, config, device='cuda'): | |
| super(DiacritizerD3, self).__init__() | |
| self.max_word_len = config["train"]["max-word-len"] | |
| self.max_sent_len = config["train"]["max-sent-len"] | |
| self.char_embed_dim = config["train"]["char-embed-dim"] | |
| self.sent_dropout_p = config["train"]["sent-dropout"] | |
| self.diac_dropout_p = config["train"]["diac-dropout"] | |
| self.vertical_dropout = config['train']['vertical-dropout'] | |
| self.recurrent_dropout = config['train']['recurrent-dropout'] | |
| self.recurrent_dropout_mode = config['train'].get('recurrent-dropout-mode', 'gal_tied') | |
| self.recurrent_activation = config['train'].get('recurrent-activation', 'sigmoid') | |
| self.sent_lstm_units = config["train"]["sent-lstm-units"] | |
| self.word_lstm_units = config["train"]["word-lstm-units"] | |
| self.decoder_units = config["train"]["decoder-units"] | |
| self.sent_lstm_layers = config["train"]["sent-lstm-layers"] | |
| self.word_lstm_layers = config["train"]["word-lstm-layers"] | |
| self.cell = config['train'].get('rnn-cell', 'lstm') | |
| self.num_layers = config["train"].get("num-layers", 2) | |
| self.RNN_Layer = K_LSTM | |
| self.batch_first = config['train'].get('batch-first', True) | |
| self.baseline = config["train"].get("baseline", False) | |
| self.device = device | |
| def build(self, wembs: T.Tensor, abjad_size: int): | |
| self.closs = F.cross_entropy | |
| self.bloss = F.binary_cross_entropy_with_logits | |
| rnn_kargs = dict( | |
| recurrent_dropout_mode=self.recurrent_dropout_mode, | |
| recurrent_activation=self.recurrent_activation, | |
| ) | |
| self.sent_lstm = self.RNN_Layer( | |
| input_size=300, | |
| hidden_size=self.sent_lstm_units, | |
| num_layers=self.sent_lstm_layers, | |
| bidirectional=True, | |
| vertical_dropout=self.vertical_dropout, | |
| recurrent_dropout=self.recurrent_dropout, | |
| batch_first=self.batch_first, | |
| **rnn_kargs, | |
| ) | |
| self.word_lstm = self.RNN_Layer( | |
| input_size=self.sent_lstm_units * 2 + self.char_embed_dim, | |
| hidden_size=self.word_lstm_units, | |
| num_layers=self.word_lstm_layers, | |
| bidirectional=True, | |
| vertical_dropout=self.vertical_dropout, | |
| recurrent_dropout=self.recurrent_dropout, | |
| batch_first=self.batch_first, | |
| return_states=True, | |
| **rnn_kargs, | |
| ) | |
| self.char_embs = nn.Embedding( | |
| abjad_size, | |
| self.char_embed_dim, | |
| padding_idx=0, | |
| ) | |
| self.attention = Attention( | |
| kind="dot", | |
| query_dim=self.word_lstm_units * 2, | |
| input_dim=self.sent_lstm_units * 2, | |
| ) | |
| self.lstm_decoder = self.RNN_Layer( | |
| input_size=self.word_lstm_units * 2 + self.attention.Dout + 8, | |
| hidden_size=self.word_lstm_units * 2, | |
| num_layers=1, | |
| bidirectional=False, | |
| vertical_dropout=self.vertical_dropout, | |
| recurrent_dropout=self.recurrent_dropout, | |
| batch_first=self.batch_first, | |
| return_states=True, | |
| **rnn_kargs, | |
| ) | |
| self.word_embs = T.tensor(wembs, dtype=T.float32) | |
| self.classifier = nn.Linear(self.lstm_decoder.hidden_size, 15) | |
| self.dropout = nn.Dropout(0.2) | |
| def forward(self, sents, words, labels): | |
| #^ sents : [b ts] | |
| #^ words : [b ts tw] | |
| #^ labels: [b ts tw] | |
| word_mask = words.ne(0.).float() | |
| #^ word_mask: [b ts tw] | |
| if self.training: | |
| q = 1.0 - self.sent_dropout_p | |
| sdo = T.bernoulli(T.full(sents.shape, q)) | |
| sents_do = sents * sdo.long() | |
| #^ sents_do : [b ts] ; DO(ts) | |
| wembs = self.word_embs[sents_do] | |
| #^ wembs : [b ts dw] ; DO(ts) | |
| else: | |
| wembs = self.word_embs[sents] | |
| #^ wembs : [b ts dw] | |
| sent_enc = self.sent_lstm(wembs.to(self.device)) | |
| #^ sent_enc : [b ts dwe] | |
| sentword_do = sent_enc.unsqueeze(2) | |
| #^ sentword_do : [b ts _ dwe] | |
| sentword_do = self.dropout(sentword_do * word_mask.unsqueeze(-1)) | |
| #^ sentword_do : [b ts tw dwe] | |
| word_index = words.view(-1, self.max_word_len) | |
| #^ word_index: [b*ts tw]? | |
| cembs = self.char_embs(word_index) | |
| #^ cembs : [b*ts tw dc] | |
| sentword_do = sentword_do.view(-1, self.max_word_len, self.sent_lstm_units * 2) | |
| #^ sentword_do : [b*ts tw dwe] | |
| char_embs = T.cat([cembs, sentword_do], dim=-1) | |
| #^ char_embs : [b*ts tw dcw] ; dcw = dc + dwe | |
| char_enc, _ = self.word_lstm(char_embs) | |
| #^ char_enc: [b*ts tw dce] | |
| char_enc_reshaped = char_enc.view(-1, self.max_sent_len, self.max_word_len, self.word_lstm_units * 2) | |
| #^ char_enc: [b ts tw dce] | |
| omit_self_mask = (1.0 - T.eye(self.max_sent_len)).unsqueeze(0).to(self.device) | |
| attn_enc, attn_map = self.attention(char_enc_reshaped, sent_enc, word_mask.bool(), prejudice_mask=omit_self_mask) | |
| #^ attn_enc: [b ts tw dae] | |
| attn_enc = attn_enc.view(-1, self.max_sent_len*self.max_word_len, self.attention.Dout) | |
| #^ attn_enc: [b*ts tw dae] | |
| if self.training and self.diac_dropout_p > 0: | |
| q = 1.0 - self.diac_dropout_p | |
| ddo = T.bernoulli(T.full(labels.shape[:-1], q)) | |
| labels = labels * ddo.unsqueeze(-1).long().to(self.device) | |
| #^ labels : [b ts tw] ; DO(ts) | |
| labels = labels.view(-1, self.max_sent_len*self.max_word_len, 8).float() | |
| #^ labels: [b*ts tw 8] | |
| char_enc = char_enc.view(-1, self.max_sent_len*self.max_word_len, self.word_lstm_units * 2) | |
| final_vec = T.cat([attn_enc, char_enc, labels], dim=-1) | |
| #^ final_vec: [b ts*tw dae+8] | |
| dec_out, _ = self.lstm_decoder(final_vec) | |
| #^ dec_out: [b*ts tw du] | |
| dec_out = dec_out.reshape(-1, self.max_word_len, self.lstm_decoder.hidden_size) | |
| diac_out = self.classifier(self.dropout(dec_out)) | |
| #^ diac_out: [b*ts tw 7] | |
| diac_out = diac_out.view(-1, self.max_sent_len, self.max_word_len, 15) | |
| #^ diac_out: [b ts tw 7] | |
| if not self.batch_first: | |
| diac_out = diac_out.swapaxes(1, 0) | |
| return diac_out, attn_map | |
| def predict_sample(self, sents, words, labels): | |
| word_mask = words.ne(0.).float() | |
| #^ mask: [b ts tw 1] | |
| if self.training: | |
| q = 1.0 - self.sent_dropout_p | |
| sdo = T.bernoulli(T.full(sents.shape, q)) | |
| sents_do = sents * sdo.long() | |
| #^ sents_do : [b ts] ; DO(ts) | |
| wembs = self.word_embs[sents_do] | |
| #^ wembs : [b ts dw] ; DO(ts) | |
| else: | |
| wembs = self.word_embs[sents] | |
| #^ wembs : [b ts dw] | |
| sent_enc = self.sent_lstm(wembs.to(self.device)) | |
| #^ sent_enc : [b ts dwe] | |
| sentword_do = sent_enc.unsqueeze(2) | |
| #^ sentword_do : [b ts _ dwe] | |
| sentword_do = self.dropout(sentword_do * word_mask.unsqueeze(-1)) | |
| #^ sentword_do : [b ts tw dwe] | |
| word_index = words.view(-1, self.max_word_len) | |
| #^ word_index: [b*ts tw]? | |
| cembs = self.char_embs(word_index) | |
| #^ cembs : [b*ts tw dc] | |
| sentword_do = sentword_do.view(-1, self.max_word_len, self.sent_lstm_units * 2) | |
| #^ sentword_do : [b*ts tw dwe] | |
| char_embs = T.cat([cembs, sentword_do], dim=-1) | |
| #^ char_embs : [b*ts tw dcw] ; dcw = dc + dwe | |
| char_enc, _ = self.word_lstm(char_embs) | |
| #^ char_enc: [b*ts tw dce] | |
| #^ word_states: ([b*ts dce], [b*ts dce]) | |
| char_enc = char_enc.view(-1, self.max_sent_len, self.max_word_len, self.word_lstm_units*2) | |
| #^ char_enc: [b ts tw dce] | |
| omit_self_mask = (1.0 - T.eye(self.max_sent_len)).unsqueeze(0).to(self.device) | |
| attn_enc, _ = self.attention(char_enc, sent_enc, word_mask.bool(), prejudice_mask=omit_self_mask) | |
| #^ attn_enc: [b ts tw dae] | |
| all_out = T.zeros(*char_enc.size()[:-1], 15).to(self.device) | |
| #^ all_out: [b ts tw 7] | |
| batch_sz = char_enc.size()[0] | |
| #^ batch_sz: b | |
| zeros = T.zeros(1, batch_sz, self.lstm_decoder.hidden_size).to(self.device) | |
| #^ zeros: [1 b du] | |
| bos_tag = T.tensor([0,0,0,0,0,1,0,0]).unsqueeze(0) | |
| #^ bos_tag: [1 8] | |
| prev_label = T.cat([bos_tag]*batch_sz).to(self.device).float() | |
| # bos_vec = T.cat([bos_tag]*batch_sz).to(self.device).float() | |
| #^ prev_label: [b 8] | |
| for ts in range(self.max_sent_len): | |
| dec_hx = (zeros, zeros) | |
| #^ dec_hx: [1 b du] | |
| for tw in range(self.max_word_len): | |
| final_vec = T.cat([attn_enc[:,ts,tw,:], char_enc[:,ts,tw,:], prev_label], dim=-1).unsqueeze(1) | |
| #^ final_vec: [b 1 dce+8] | |
| dec_out, dec_hx = self.lstm_decoder(final_vec, dec_hx) | |
| #^ dec_out: [b 1 du] | |
| dec_out = dec_out.squeeze(0) | |
| dec_out = dec_out.transpose(0,1) | |
| logits_raw = self.classifier(self.dropout(dec_out)) | |
| #^ logits_raw: [b 1 15] | |
| out_idx = T.max(T.softmax(logits_raw.squeeze(), dim=-1), dim=-1)[1] | |
| haraka, tanween, shadda = flat2_3head(out_idx.detach().cpu().numpy()) | |
| haraka_onehot = T.eye(6)[haraka].float().to(self.device) | |
| #^ haraka_onehot+bos_tag: [b 6] | |
| tanween = T.tensor(tanween).float().unsqueeze(-1).to(self.device) | |
| shadda = T.tensor(shadda).float().unsqueeze(-1).to(self.device) | |
| prev_label = T.cat([haraka_onehot, tanween, shadda], dim=-1) | |
| all_out[:,ts,tw,:] = logits_raw.squeeze() | |
| if not self.batch_first: | |
| all_out = all_out.swapaxes(1, 0) | |
| return all_out | |
| def step(self, xt, yt, mask=None): | |
| xt[1] = xt[1].to(self.device) | |
| xt[2] = xt[2].to(self.device) | |
| #^ yt: [b ts tw] | |
| yt = yt.to(self.device) | |
| if self.training: | |
| diac, _ = self(*xt) | |
| else: | |
| diac = self.predict_sample(*xt) | |
| #^ diac[0] : [b ts tw 5] | |
| loss = self.closs(diac.view(-1,15), yt.view(-1)) | |
| return loss | |
| def predict(self, dataloader): | |
| training = self.training | |
| self.eval() | |
| preds = {'haraka': [], 'shadda': [], 'tanween': []} | |
| print("> Predicting...") | |
| for inputs, _ in tqdm(dataloader, total=len(dataloader)): | |
| inputs[1] = inputs[1].to(self.device) | |
| inputs[2] = inputs[2].to(self.device) | |
| diac = self.predict_sample(*inputs) | |
| output = np.argmax(T.softmax(diac.detach(), dim=-1).cpu().numpy(), axis=-1) | |
| #^ [b ts tw] | |
| haraka, tanween, shadda = flat_2_3head(output) | |
| preds['haraka'].extend(haraka) | |
| preds['tanween'].extend(tanween) | |
| preds['shadda'].extend(shadda) | |
| self.train(training) | |
| return ( | |
| np.array(preds['haraka']), | |
| np.array(preds["tanween"]), | |
| np.array(preds["shadda"]), | |
| ) | |
| if __name__ == "__main__": | |
| import yaml | |
| config_path = "configs/dd/config_d2.yaml" | |
| model_path = "models/tashkeela-d2.pt" | |
| with open(config_path, 'r', encoding="utf-8") as file: | |
| config = yaml.load(file, Loader=yaml.FullLoader) | |
| data_utils = DatasetUtils(config) | |
| vocab_size = len(data_utils.letter_list) | |
| word_embeddings = data_utils.embeddings | |
| model = DiacritizerD2(config, device='cpu') | |
| model.build(word_embeddings, vocab_size) | |
| model.load_state_dict(T.load(model_path, map_location=T.device('cpu'))["state_dict"]) |