mininfradev/seq2seq-kisl-gloss-space

1import json, re, unicodedata, torch, torch.nn as nn, torch.nn.functional as F2 3device = torch.device("cuda" if torch.cuda.is_available() else "cpu")4SOS_token, EOS_token = 0, 15MAX_LENGTH = 206 7def unicodeToAscii(s):8    return ''.join(c for c in unicodedata.normalize('NFD', s) if unicodedata.category(c) != 'Mn')9 10def normalizeString(s):11    s = unicodeToAscii(s.lower().strip())12    s = re.sub(r"([.!?])", r" \1", s)13    s = re.sub(r"[^a-zA-Z!?]+", r" ", s)14    return s.strip()15 16class EncoderRNN(nn.Module):17    def __init__(self, input_size, hidden_size, dropout_p=0.1):18        super().__init__()19        self.embedding = nn.Embedding(input_size, hidden_size)20        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True)21        self.dropout = nn.Dropout(dropout_p)22    def forward(self, inp):23        embedded = self.dropout(self.embedding(inp))24        output, hidden = self.gru(embedded)25        return output, hidden26 27class BahdanauAttention(nn.Module):28    def __init__(self, hidden_size):29        super().__init__()30        self.Wa = nn.Linear(hidden_size, hidden_size)31        self.Ua = nn.Linear(hidden_size, hidden_size)32        self.Va = nn.Linear(hidden_size, 1)33    def forward(self, query, keys):34        scores = self.Va(torch.tanh(self.Wa(query) + self.Ua(keys)))  # (B,T,1)35        scores = scores.squeeze(2).unsqueeze(1)                        # (B,1,T)36        weights = F.softmax(scores, dim=-1)                            # (B,1,T)37        context = torch.bmm(weights, keys)                             # (B,1,H)38        return context, weights39 40class AttnDecoderRNN(nn.Module):41    def __init__(self, hidden_size, output_size, dropout_p=0.1):42        super().__init__()43        self.embedding = nn.Embedding(output_size, hidden_size)44        self.attention = BahdanauAttention(hidden_size)45        self.gru = nn.GRU(2*hidden_size, hidden_size, batch_first=True)46        self.out = nn.Linear(hidden_size, output_size)47        self.dropout = nn.Dropout(dropout_p)48 49    def forward_step(self, inp, hidden, encoder_outputs):50        embedded = self.dropout(self.embedding(inp))                   # (B,1,H)51        query = hidden.permute(1,0,2)                                  # (B,1,H)52        context, attn_weights = self.attention(query, encoder_outputs) # (B,1,H),(B,1,T)53        x = torch.cat((embedded, context), dim=2)                      # (B,1,2H)54        output, hidden = self.gru(x, hidden)                           # (B,1,H),(1,B,H)55        output = self.out(output)                                      # (B,1,V)56        return output, hidden, attn_weights57 58    def forward(self, encoder_outputs, encoder_hidden, target_tensor=None):59        B = encoder_outputs.size(0)60        decoder_input = torch.empty(B, 1, dtype=torch.long, device=device).fill_(SOS_token)61        decoder_hidden = encoder_hidden62        outs, atts = [], []63        for _ in range(MAX_LENGTH):64            o, decoder_hidden, aw = self.forward_step(decoder_input, decoder_hidden, encoder_outputs)65            outs.append(o)66            atts.append(aw)67            if target_tensor is not None:68                decoder_input = target_tensor[:, _].unsqueeze(1)69            else:70                _, topi = o.topk(1)                    # (B,1,1)71                decoder_input = topi.squeeze(-1)       # (B,1)72        outs = torch.cat(outs, dim=1)                   # (B,T,V)73        outs = F.log_softmax(outs, dim=-1)74        atts = torch.cat(atts, dim=1)                   # (B,T,EncT)75        return outs, decoder_hidden, atts76 77def load_vocab(path):78    d = json.load(open(path, "r", encoding="utf-8"))79    # Re-cast keys for index2word back to ints80    d["index2word"] = {int(k): v for k, v in d["index2word"].items()}81    return d82 83def sentence_to_tensor(lang_dict, sentence):84    # NOTE: notebook has no UNK/PAD; drop OOV tokens to avoid KeyError at inference85    w2i = lang_dict["word2index"]86    toks = [w for w in normalizeString(sentence).split() if w in w2i]87    idxs = [w2i[w] for w in toks] + [EOS_token]88    idxs = idxs[:MAX_LENGTH]89    return torch.tensor([idxs], dtype=torch.long, device=device)      # (1,L)90 91@torch.no_grad()92def translate_text(text, input_lang, output_lang, encoder, decoder):93    src = sentence_to_tensor(input_lang, text)94    enc_outs, enc_hid = encoder(src)                                  # (1,T,H),(1,1,H)95    dec_outs, _, _ = decoder(enc_outs, enc_hid, target_tensor=None)   # (1,MAX,V)96    _, topi = dec_outs.topk(1)                                        # (1,MAX,1)97    ids = topi.squeeze(-1).squeeze(0).tolist()98 99    i2w = output_lang["index2word"]100    words = []101    for idx in ids:102        if idx == EOS_token:          # stop decoding103            break104        if idx in {SOS_token, EOS_token}:  # skip special tokens105            continue106        words.append(i2w.get(idx, "<UNK>"))107 108    return " ".join(words)109 110