රෙට්රෝආකෘතිය

16import math
17from typing import Set
18
19import torch
20from torch import nn
21
22from labml.logger import inspect

25class RotaryPositionalEmbeddings(nn.Module):

36    def __init__(self, d: int, base: int = 10_000):

41        super().__init__()

43        self.theta = nn.Parameter(1. / (base ** (torch.arange(0, d, 2).float() / d)), requires_grad=False)

45    def forward(self, x: torch.Tensor):

50        batch_size, seq_len, n_heads, d = x.shape

53        d_2 = d // 2

56        seq_idx = torch.arange(seq_len, device=x.device).type_as(self.theta)

59        idx_theta = torch.einsum('n,d->nd', seq_idx, self.theta)

63        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)

67        neg_half_x = torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)

79        rx = (x * idx_theta2.cos()[None, :, None, :]) + (neg_half_x * idx_theta2.sin()[None, :, None, :])

82        return rx

85class SelfAttention(nn.Module):

92    def __init__(self, d_model: int, n_heads: int, d_k: int, is_causal: bool):

99        super().__init__()
100
101        self.is_causal = is_causal
102        self.n_heads = n_heads
103        self.d_k = d_k

106        self.scale = 1 / math.sqrt(self.d_k)

109        self.query = nn.Linear(d_model, n_heads * d_k)
110        self.key = nn.Linear(d_model, n_heads * d_k)
111        self.value = nn.Linear(d_model, n_heads * d_k)

114        self.norm = nn.LayerNorm(d_model)

117        self.softmax = nn.Softmax(dim=-1)

120        self.rotary_pe = RotaryPositionalEmbeddings(self.d_k)

123        self.output = nn.Linear(n_heads * d_k, d_model)

125    def mask_attention(self, attn: torch.Tensor):

133        if not self.is_causal:
134            return attn

137        mask = torch.tril(attn.new_ones(attn.shape[-2:]))

139        return attn.masked_fill(mask == 0, float('-inf'))

141    def forward(self, h: torch.Tensor):

147        h_res = h

150        h = self.norm(h)

154        mh_shape = (*h.shape[:-1], self.n_heads, self.d_k)
155        q = self.query(h).view(mh_shape)
156        k = self.key(h).view(mh_shape)
157        v = self.value(h).view(mh_shape)

160        q = self.rotary_pe(q)
161        k = self.rotary_pe(k)

164        attn = torch.einsum('bihd,bjhd->bhij', q, k)

166        attn = attn * self.scale

169        attn = self.mask_attention(attn)

172        attn = self.softmax(attn)

175        h = torch.einsum("bhij,bjhd->bihd", attn, v)

179        h = h.reshape(*h.shape[:-2], -1)

183        h = self.output(h)

186        return h + h_res

189class CrossAttention(nn.Module):

203    def __init__(self, d_model: int, n_heads: int, d_k: int):

209        super().__init__()
210
211        self.n_heads = n_heads
212        self.d_k = d_k

215        self.scale = 1 / math.sqrt(self.d_k)

218        self.query = nn.Linear(d_model, n_heads * d_k)
219        self.key = nn.Linear(d_model, n_heads * d_k)
220        self.value = nn.Linear(d_model, n_heads * d_k)

223        self.norm = nn.LayerNorm(d_model)

226        self.softmax = nn.Softmax(dim=-1)

229        self.output = nn.Linear(n_heads * d_k, d_model)

231    def forward(self, e: torch.Tensor, h: torch.Tensor):

240        e_res = e

243        e = self.norm(e)

246        q = self.query(e).view(*e.shape[:-1], self.n_heads, self.d_k)

248        k = self.key(h).view(*h.shape[:-1], self.n_heads, self.d_k)
249        v = self.value(h).view(*h.shape[:-1], self.n_heads, self.d_k)

254        attn = torch.einsum('bcnihd,bcjhd->bcnhij', q, k)

256        attn = attn * self.scale

259        attn = self.softmax(attn)

262        e = torch.einsum("bcnhij,bcjhd->bcnihd", attn, v)

266        e = e.reshape(*e.shape[:-2], -1)

270        e = self.output(e)

273        return e + e_res

276class ChunkedCrossAttention(nn.Module):

288    def __init__(self, d_model: int, n_heads: int, d_k: int, chunk_len: int):

296        super().__init__()
297
298        self.chunk_len = chunk_len
299        self.n_heads = n_heads
300        self.d_k = d_k

303        self.scale = 1 / math.sqrt(self.d_k)

306        self.query = nn.Linear(d_model, n_heads * d_k)
307        self.key = nn.Linear(d_model, n_heads * d_k)
308        self.value = nn.Linear(d_model, n_heads * d_k)

311        self.norm = nn.LayerNorm(d_model)

314        self.softmax = nn.Softmax(dim=-1)

317        self.output = nn.Linear(n_heads * d_k, d_model)

319    def forward(self, h: torch.Tensor, e: torch.Tensor):

326        batch_size, chunks, neighbors, neighbor_len, d_model = e.shape

329        if chunks == 0:
330            return h

333        h_res = h

341        h = h[:, self.chunk_len - 1:]

343        h = self.norm(h)

345        if h.shape[1] < chunks * self.chunk_len:
346            h = torch.cat((h, h.new_zeros(batch_size, chunks * self.chunk_len - h.shape[1], d_model)), dim=1)

348        h = h.reshape(batch_size, chunks, self.chunk_len, d_model)

351        q = self.query(h).view(*h.shape[:-1], self.n_heads, self.d_k)

353        k = self.key(e).view(*e.shape[:-1], self.n_heads, self.d_k)
354        v = self.value(e).view(*e.shape[:-1], self.n_heads, self.d_k)

359        attn = torch.einsum('bcihd,bcnjhd->bchinj', q, k)

361        attn = attn * self.scale

364        attn = self.softmax(attn.view(*attn.shape[:-2], -1)).view(attn.shape)

367        h = torch.einsum("bchinj,bcnjhd->bcihd", attn, v)

371        h = h.reshape(batch_size, chunks * self.chunk_len, -1)

375        h = self.output(h)

378        h = torch.cat((h.new_zeros(batch_size, self.chunk_len - 1, d_model), h), dim=1)

381        return h[:, :h_res.shape[1]] + h_res

384class FeedForward(nn.Module):

391    def __init__(self, d_model: int, d_ff: int):

397        super().__init__()

400        self.lin1 = nn.Linear(d_model, d_ff)
401        self.lin2 = nn.Linear(d_ff, d_model)

404        self.act = nn.ReLU()

407        self.norm = nn.LayerNorm(d_model)

409    def forward(self, h: torch.Tensor):

415        h_res = h

417        h = self.norm(h)

419        h = self.lin1(h)

421        h = self.act(h)

423        h = self.lin2(h)

426        return h + h_res

429class NearestNeighborEncoder(nn.Module):

436    def __init__(self, chunk_len: int, n_layers: int, ca_layers: Set[int],
437                 d_model: int, n_heads: int, d_k: int, d_ff: int):

448        super().__init__()
449        self.ca_layers = ca_layers
450        self.chunk_len = chunk_len

452        self.ca = nn.ModuleList([CrossAttention(d_model, n_heads, d_k) for _ in range(len(ca_layers))])

454        self.attn = nn.ModuleList([SelfAttention(d_model, n_heads, d_k, is_causal=False) for _ in range(n_layers)])

456        self.ffw = nn.ModuleList([FeedForward(d_model, d_ff) for _ in range(n_layers)])

459        self.norm_h = nn.LayerNorm(d_model)

461    def forward(self, e: torch.Tensor, h: torch.Tensor):

474        batch_size, chunks, neighbors, neighbor_len, d_model = e.shape

477        h_split = h[:, :self.chunk_len * chunks, :].reshape(batch_size, chunks, self.chunk_len, d_model)

480        h_split = self.norm_h(h_split)

483        p_ca = 0

485        for p in range(len(self.attn)):

488            e = self.attn[p](e.view(-1, neighbor_len, d_model)).view(e.shape)

491            if p in self.ca_layers:

493                e = self.ca[p_ca](e, h_split)

495                p_ca += 1

498            e = self.ffw[p](e)

501        return e

504class RetroModel(nn.Module):

511    def __init__(self, n_vocab: int, d_model: int, n_layers: int, ca_layers: Set[int], chunk_len: int,
512                 n_heads: int, d_k: int, d_ff: int, encoder: NearestNeighborEncoder):

524        super().__init__()
525
526        self.ca_layers = ca_layers
527        self.encoder = encoder

530        self.emb = nn.Embedding(n_vocab, d_model)

532        self.cca = nn.ModuleList(
533            [ChunkedCrossAttention(d_model, n_heads, d_k, chunk_len) for _ in range(len(ca_layers))])

535        self.attn = nn.ModuleList([SelfAttention(d_model, n_heads, d_k, is_causal=True) for _ in range(n_layers)])

537        self.ffw = nn.ModuleList([FeedForward(d_model, d_ff) for _ in range(n_layers)])

539        self.read = nn.Linear(d_model, n_vocab)

543        self.norm_e = nn.LayerNorm(d_model)

545    def forward(self, x: torch.Tensor, ret: torch.Tensor):

554        h = self.emb(x)

560        ret_emb = self.emb(ret)

563        p_ca = 0

565        for p in range(len(self.attn)):

567            h = self.attn[p](h)

571            if self.ca_layers and p == min(self.ca_layers):

575                e = self.encoder(ret_emb, h)

577                e = self.norm_e(e)

580            if p in self.ca_layers:

582                h = self.cca[p_ca](h, e)

584                p_ca += 1

587            h = self.ffw[p](h)

590        return self.read(h)

593def _test():

597    chunk_len = 4
598    d_model = 8
599    d_ff = 32
600    n_heads = 2
601    d_k = 4
602
603    device = torch.device('cuda:0')
604
605    m = RetroModel(5, d_model, 6, {2, 5}, chunk_len, n_heads, d_k, d_ff,
606                   encoder=NearestNeighborEncoder(chunk_len, 2, {1}, d_model, n_heads, d_k, d_ff))
607
608    m.to(device)
609    x = [1, 2, 4, 4, 0, 1, 2, 3, 4, 3]
610    ret = [
611        [[0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1]],
612        [[0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1]],
613    ]
614    res = m(torch.tensor([x] * 10).to(device), torch.tensor([ret] * 10).to(device))
615
616    inspect(res)

620if __name__ == '__main__':
621    _test()