Create model.py

model.py

@@ -0,0 +1,158 @@
+from typing import Sequence
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.models.vision_transformer import PatchEmbed, VisionTransformer
+from dataclasses import dataclass
+from torch import Tensor
+import math
+
+
+class ImageEncoder(VisionTransformer):
+    def __init__(self, config):
+        super().__init__(
+            img_size=config.img_size,
+            patch_size=config.patch_size,
+            in_chans=config.n_channel,
+            embed_dim=config.n_embed,
+            depth=config.n_layer,
+            num_heads=config.n_head,
+            mlp_ratio=4,
+            qkv_bias=True,
+            drop_rate=0.0,
+            attn_drop_rate=0.0,
+            drop_path_rate=0.0,
+            embed_layer=PatchEmbed,
+            num_classes=0,  # These
+            global_pool='',  # disable the
+            class_token=False,  # classifier head.
+        )
+
+    def forward(self, x):
+        return self.forward_features(x)
+
+
+class RMSNorm(nn.RMSNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+class Linear(nn.Linear):
+    def forward(self, x: Tensor) -> Tensor:
+        return F.linear(x, self.weight.to(x.dtype), None if self.bias is None else self.bias.to(x.dtype))
+
+
+class TextDecoder(nn.Module):
+    def __init__(self, config, ) -> None:
+        super().__init__()
+        self.config = config
+        self.n_head = 2 * config.n_head
+        self.tok_embed = nn.Embedding(config.vocab_size, config.n_embed)
+        self.pos_embed = nn.Parameter(torch.Tensor(
+            1, config.block_size, config.n_embed))
+        self.dropout = nn.Dropout(config.dropout)
+
+        self.sa_ln = RMSNorm(config.n_embed)
+        self.sa_attn = nn.MultiheadAttention(config.n_embed, self.n_head, dropout=config.dropout, batch_first=True)
+
+        self.cross_ln = RMSNorm(config.n_embed)
+        self.cross_attn = nn.MultiheadAttention(config.n_embed, self.n_head, dropout=config.dropout, batch_first=True)
+
+        self.ffn_ln = RMSNorm(config.n_embed)
+        dim_feedforward = 4*config.n_embed
+        self.ffn = nn.Sequential(
+            Linear(config.n_embed, dim_feedforward, bias=config.bias),
+            nn.GELU(),
+            Linear(dim_feedforward, config.n_embed, bias=config.bias),
+            nn.Dropout(config.dropout)
+        )
+        self.lm_head = Linear(config.n_embed, config.vocab_size)
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+
+    def forward(self, x: Tensor, xi: Tensor):
+        """
+        x: input token ids
+        xi: image features (already normalized by ImageEncoder)
+        """
+        b, t = x.size()
+        tok_embed = self.tok_embed(x) * math.sqrt(self.config.n_embed)
+
+        ctx = torch.cat(
+            [tok_embed[:, :1], self.pos_embed[:, :t-1] + tok_embed[:, 1:]], dim=1)
+        ctx = self.dropout(ctx)
+        ctx = self.sa_ln(ctx)
+        res = self.dropout(self.pos_embed[:, :t].expand(b, -1, -1))  # (b, t, n_embed)
+
+        mask = torch.triu(torch.ones((t, t), dtype=torch.bool, device=x.device), 1)
+        query, sa_weights = self.sa_attn(self.sa_ln(res), ctx, ctx, attn_mask=mask)
+        res = res + query
+        query, ca_weights = self.cross_attn(self.cross_ln(res), xi, xi)
+        res = res + query
+        res = res + self.ffn(self.ffn_ln(res))
+        return self.lm_head(res[:, [-1], :]).float()
+
+
+class OCRModel(nn.Module):
+    def __init__(self, config, tokenizer) -> None:
+        super().__init__()
+        self.encoder = ImageEncoder(config)
+        self.decoder = TextDecoder(config)
+        self.tokenizer = tokenizer
+
+    def forward(self, img_tensor: Tensor, input_tokens: Tensor):
+        xi = self.encoder(img_tensor)
+        logits, loss = self.decoder(input_tokens, xi)
+        return logits, loss
+
+    @torch.inference_mode()
+    def generate(self, img_tensor: Tensor, max_new_tokens: int, temperature=1.0, top_k=None):
+        xi = self.encoder(img_tensor.unsqueeze(0))
+        idx = torch.full((xi.size(0),1), fill_value=self.tokenizer.bos_id, dtype=torch.long, device=img_tensor.device)
+        for i in range(max_new_tokens):
+            logits = self.decoder(idx, xi)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+            if idx_next.item() == self.tokenizer.eos_id:
+                break
+        return self.tokenizer.decode(idx[0].tolist(), ignore_special_tokens=True)
+
+
+@dataclass
+class ModelConfig:
+    img_size: Sequence[int]
+    patch_size: Sequence[int]
+    n_channel: int
+    vocab_size: int
+    block_size: int
+    n_layer: int
+    n_head: int
+    n_embed: int
+    dropout: float = 0.0
+    bias: bool = True
+
+
+def load_model():
+    import pickle
+    with open('tokenizer.pkl', 'rb') as inp:
+        tokenizer = pickle.load(inp)
+    config = ModelConfig(
+        img_size=(32, 128),
+        patch_size=(4, 8),
+        n_channel=3,
+        vocab_size=len(tokenizer),
+        block_size=192,
+        n_layer=12,
+        n_head=3,
+        n_embed=192,
+        dropout=0.1,
+        bias=True,
+    )
+    model = OCRModel(config, tokenizer)
+    state_dict = torch.hub.load_state_dict_from_url('https://huggingface.co/KrorngAI/PARSeqForKhmer/resolve/main/parseq_kh.pt', map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict)
+    return model