Use tror-yong-ocr package

model.py

@@ -1,206 +1,22 @@
-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 CharTokenizer:
-    def __init__(self, chars, special_tokens=['<s>', '</s>', '<pad>', '<unk>']):
-        self.special_tokens = special_tokens
-        # Unique characters + special tokens
-        self.vocab = tuple(special_tokens[:1]) + tuple(chars) + tuple(special_tokens[1:])
-        self.str_to_int = {s: i for i, s in enumerate(self.vocab)}
-        self.int_to_str = {i: s for i, s in enumerate(self.vocab)}
-        self.bos_id = self.str_to_int['<s>']
-        self.eos_id = self.str_to_int['</s>']
-        self.pad_id = self.str_to_int['<pad>']
-        self.unk_id = self.str_to_int['<unk>']
-
-    def __len__(self):
-        return len(self.vocab)
-
-    def encode(self, text, add_special_tokens=False):
-        tokens = []
-        i = 0
-        while i < len(text):
-            matched_special = False
-            # Check for existing special tokens in the input string
-            for spec in self.special_tokens:
-                if text.startswith(spec, i):
-                    tokens.append(self.str_to_int[spec])
-                    i += len(spec)
-                    matched_special = True
-                    break
-
-            if not matched_special:
-                char = text[i]
-                tokens.append(self.str_to_int.get(char, self.str_to_int['<unk>']))
-                i += 1
-
-        # Wrap with <s> and </s> if requested
-        if add_special_tokens:
-            tokens = [self.str_to_int['<s>']] + tokens + [self.str_to_int['</s>']]
-
-        return tokens
-
-    def decode(self, ids, ignore_special_tokens=False):
-        if ignore_special_tokens:
-            # Filter out any ID that belongs to the special_tokens list
-            return "".join([self.int_to_str[i] for i in ids if self.int_to_str[i] not in self.special_tokens])
-
-        return "".join([self.int_to_str.get(i, '<unk>') for i in ids])
-
-
-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
+from tror_yong_ocr import TrorYongOCR, TrorYongConfig
+from tror_yong_ocr import get_tokenizer
 
 
 def load_model():
-    kh_charset = "០១២៣៤៥៦៧៨៩កខគឃងចឆជឈញដឋឌឍណតថទធនបផពភមយរលវសហឡអឥឧឳឪឱឫឬឭឮឦឰឯាិីឹឺុូួើឿៀេែៃោៅំះៈ់៉៊៍័៏៌្ ។៕៖ៗ"
-    en_charset = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
-    tokenizer = CharTokenizer(en_charset+kh_charset)
-    
-    config = ModelConfig(
+    tokenizer = get_tokenizer()
+    config = TrorYongConfig(
         img_size=(32, 128),
         patch_size=(4, 8),
         n_channel=3,
-        vocab_size=len(tokenizer),
+        vocab_size=len(tokenizer), # exclude pad and unk tokens
         block_size=192,
-        n_layer=12,
-        n_head=3,
-        n_embed=192,
+        n_layer=4,
+        n_head=6,
+        n_embed=384,
         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 = TrorYongOCR(config, tokenizer)
+    state_dict = torch.hub.load_state_dict_from_url('https://huggingface.co/KrorngAI/PARSeqForKhmer/resolve/main/best_model-80epoch.pt', map_location=torch.device('cpu'))
     model.load_state_dict(state_dict)
     return model