Update handler.py

handler.py

@@ -2,6 +2,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 TOKENIZER_NAME = "polyglots/Extended-Sinhala-LLaMA"
+VOCAB_SIZE     = 139336
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -15,18 +16,28 @@ class EndpointHandler:
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
+        # ── Fix: patch vocab size in config BEFORE model is created ──────────
+        # Without this, model is built with 128256 vocab then fails to load
+        # the 139336-vocab checkpoint weights
+        print(f"Patching config vocab_size to {VOCAB_SIZE:,}...")
+        config = AutoConfig.from_pretrained(path, trust_remote_code=True)
+        config.vocab_size = VOCAB_SIZE
+
         print(f"Loading model from {path}...")
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
+            config                = config,
             torch_dtype  = torch.float16,
             device_map   = "auto",
             trust_remote_code = True,
+            ignore_mismatched_sizes = True, 
         )
         # Resize to match extended vocab (139,336 tokens)
-        self.model.resize_token_embeddings(len(self.tokenizer))
+        # self.model.resize_token_embeddings(len(self.tokenizer))
+        
         self.model.config.pad_token_id = self.tokenizer.eos_token_id
         self.model.eval()
-        print("Model ready!")
+        print(f"Ready! Vocab: {self.model.config.vocab_size:,}")
 
     def __call__(self, data: dict) -> dict:
         # ── unpack request ───────────────────────────────────────────────────