Update handler.py

handler.py

@@ -21,26 +21,34 @@ class EndpointHandler:
         # the 139336-vocab checkpoint weights
         print(f"Patching config vocab_size to {VOCAB_SIZE:,}...")
         config = LlamaConfig.from_pretrained(path)
+
+        # Force correct vocab size BEFORE model is built
+        # so embeddings are initialized at the right size
         config.vocab_size = VOCAB_SIZE
 
-        print(f"Loading model from {path}...")
-        self.model = LlamaForCausalLM.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))
-        
+        print(f"Loading model weights...")
+        self.model = LlamaForCausalLM(config)   # ← build empty model at correct size first
+
+        # Now load the checkpoint weights — sizes will match
+        import os
+        from safetensors.torch import load_file
+
+        weights = {}
+        for f in sorted(os.listdir(path)):
+            if f.endswith(".safetensors"):
+                print(f"  Loading shard: {f}")
+                weights.update(load_file(os.path.join(path, f)))
+
+        missing, unexpected = self.model.load_state_dict(weights, strict=False)
+        print(f"  Missing keys:    {len(missing)}")
+        print(f"  Unexpected keys: {len(unexpected)}")
+
+        self.model = self.model.to(torch.float16).to("cuda")
         self.model.config.pad_token_id = self.tokenizer.eos_token_id
         self.model.eval()
         print(f"Ready! Vocab: {self.model.config.vocab_size:,}")
 
-    def __call__(self, data: dict) -> dict:
-        # ── unpack request ───────────────────────────────────────────────────
+        def __call__(self, data: dict) -> dict:
         inputs      = data.get("inputs", "")
         params      = data.get("parameters", {})
         max_tokens  = params.get("max_new_tokens", 400)
@@ -49,30 +57,26 @@ class EndpointHandler:
         rep_penalty = params.get("repetition_penalty", 1.1)
 
         if not inputs:
-            return {"error": "No input text provided. Use the 'inputs' key."}
+            return {"error": "No input provided. Use the 'inputs' key."}
 
-        # ── tokenise ─────────────────────────────────────────────────────────
         tokenized = self.tokenizer(
             inputs,
-            return_tensors  = "pt",
-            truncation      = True,
-            max_length      = 1024,
+            return_tensors = "pt",
+            truncation     = True,
+            max_length     = 1024,
         ).to(self.model.device)
 
-        # ── generate ─────────────────────────────────────────────────────────
         with torch.no_grad():
             output_ids = self.model.generate(
                 **tokenized,
-                max_new_tokens       = max_tokens,
-                temperature          = temperature,
-                top_p                = top_p,
-                repetition_penalty   = rep_penalty,
-                do_sample            = True,
-                pad_token_id         = self.tokenizer.eos_token_id,
+                max_new_tokens     = max_tokens,
+                temperature        = temperature,
+                top_p              = top_p,
+                repetition_penalty = rep_penalty,
+                do_sample          = True,
+                pad_token_id       = self.tokenizer.eos_token_id,
             )
 
-        # ── decode (strip prompt, return only new tokens) ─────────────────────
         new_tokens = output_ids[0][tokenized.input_ids.shape[1]:]
         decoded    = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
-
-        return {"generated_text": decoded.strip()}
+        return {"generated_text": decoded.strip()}