Image (768×768)
↓
Florence-2 Vision Encoder (frozen, 360M params)
↓
Vision Projection (1024 → 512 dim)
↓
Transformer Decoder (4 layers, 8 heads)
↓
Character-level predictions (187 vocab)The 5% gap is expected when adapting a general vision-language model versus training a specialized OCR architecture. However, AssameseOCR offers:
1 import torch
2 import torch . nn as nn
3 from PIL import Image
4 from transformers import AutoModelForCausalLM , CLIPImageProcessor
5 from huggingface_hub import hf_hub_download
6 import json
7
8 # CharTokenizer class
9 class CharTokenizer :
10 def __init__ ( self , vocab ) :
11 self . vocab = vocab
12 self . char2id = { c : i for i , c in enumerate ( vocab ) }
13 self . id2char = { i : c for i , c in enumerate ( vocab ) }
14 self . pad_token_id = self . char2id [ "<pad>" ]
15 self . bos_token_id = self . char2id [ "<s>" ]
16 self . eos_token_id = self . char2id [ "</s>" ]
17
18 def encode ( self , text , max_length = None , add_special_tokens = True ) :
19 ids = [ self . bos_token_id ] if add_special_tokens else [ ]
20 for ch in text :
21 ids . append ( self . char2id . get ( ch , self . char2id [ "<unk>" ] ) )
22 if add_special_tokens :
23 ids . append ( self . eos_token_id )
24 if max_length :
25 ids = ids [ : max_length ]
26 if len ( ids ) < max_length :
27 ids += [ self . pad_token_id ] * ( max_length - len ( ids ) )
28 return ids
29
30 def decode ( self , ids , skip_special_tokens = True ) :
31 chars = [ ]
32 for i in ids :
33 ch = self . id2char . get ( i , "" )
34 if skip_special_tokens and ch . startswith ( "<" ) :
35 continue
36 chars . append ( ch )
37 return "" . join ( chars )
38
39 @classmethod
40 def load ( cls , path ) :
41 with open ( path , "r" , encoding = "utf-8" ) as f :
42 vocab = json . load ( f )
43 return cls ( vocab )
44
45 # FlorenceCharOCR model class
46 class FlorenceCharOCR ( nn . Module ) :
47 def __init__ ( self , florence_model , vocab_size , vision_hidden_dim , decoder_hidden_dim = 512 , num_layers = 4 ) :
48 super ( ) . __init__ ( )
49 self . florence_model = florence_model
50
51 for param in self . florence_model . parameters ( ) :
52 param . requires_grad = False
53
54 self . vision_proj = nn . Linear ( vision_hidden_dim , decoder_hidden_dim )
55 self . embedding = nn . Embedding ( vocab_size , decoder_hidden_dim )
56 decoder_layer = nn . TransformerDecoderLayer (
57 d_model = decoder_hidden_dim ,
58 nhead = 8 ,
59 batch_first = True
60 )
61 self . decoder = nn . TransformerDecoder ( decoder_layer , num_layers = num_layers )
62 self . fc_out = nn . Linear ( decoder_hidden_dim , vocab_size )
63
64 def forward ( self , pixel_values , tgt_ids , tgt_mask = None ) :
65 with torch . no_grad ( ) :
66 vision_feats = self . florence_model . _encode_image ( pixel_values )
67
68 vision_feats = self . vision_proj ( vision_feats )
69 tgt_emb = self . embedding ( tgt_ids )
70 decoder_out = self . decoder ( tgt_emb , vision_feats , tgt_mask = tgt_mask )
71 logits = self . fc_out ( decoder_out )
72
73 return logits
74
75 # Load components
76 device = "cuda" if torch . cuda . is_available ( ) else "cpu"
77
78 # Download files from HuggingFace
79 tokenizer_path = hf_hub_download ( repo_id = "MWirelabs/assamese-ocr" , filename = "assamese_char_tokenizer.json" )
80 model_path = hf_hub_download ( repo_id = "MWirelabs/assamese-ocr" , filename = "assamese_ocr_best.pt" )
81
82 # Load tokenizer
83 char_tokenizer = CharTokenizer . load ( tokenizer_path )
84
85 # Load Florence base model
86 florence_model = AutoModelForCausalLM . from_pretrained (
87 "microsoft/Florence-2-large-ft" ,
88 trust_remote_code = True
89 ) . to ( device )
90
91 # Load image processor
92 image_processor = CLIPImageProcessor . from_pretrained ( "microsoft/Florence-2-large-ft" )
93
94 # Initialize OCR model
95 ocr_model = FlorenceCharOCR (
96 florence_model = florence_model ,
97 vocab_size = len ( char_tokenizer . vocab ) ,
98 vision_hidden_dim = 1024 ,
99 decoder_hidden_dim = 512 ,
100 num_layers = 4
101 ) . to ( device )
102
103 # Load trained weights
104 checkpoint = torch . load ( model_path , map_location = device )
105 ocr_model . load_state_dict ( checkpoint [ 'model_state_dict' ] )
106 ocr_model . eval ( )
107
108 # Inference function
109 def recognize_text ( image_path ) :
110 # Load and process image
111 image = Image . open ( image_path ) . convert ( "RGB" )
112 pixel_values = image_processor ( images = [ image ] , return_tensors = "pt" ) [ 'pixel_values' ] . to ( device )
113
114 # Generate prediction
115 with torch . no_grad ( ) :
116 # Start with BOS token
117 generated_ids = [ char_tokenizer . bos_token_id ]
118
119 for _ in range ( 128 ) : # max length
120 tgt_tensor = torch . tensor ( [ generated_ids ] , device = device )
121 logits = ocr_model ( pixel_values , tgt_tensor )
122
123 # Get next token
124 next_token = logits [ 0 , - 1 ] . argmax ( ) . item ( )
125 generated_ids . append ( next_token )
126
127 # Stop if EOS
128 if next_token == char_tokenizer . eos_token_id :
129 break
130
131 # Decode
132 text = char_tokenizer . decode ( generated_ids , skip_special_tokens = True )
133 return text
134
135 # Example usage
136 result = recognize_text ( "assamese_text.jpg" )
137 print ( f"Recognized text: { result } " )
1 @software{assameseocr2026,
2 author = {MWire Labs},
3 title = {AssameseOCR: Vision-Language Model for Assamese Text Recognition},
4 year = {2026},
5 publisher = {Hugging Face},
6 url = {https://huggingface.co/MWirelabs/assamese-ocr}
7 }
