Update handler.py

handler.py

@@ -20,16 +20,16 @@ class EndpointHandler:
         # 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 = LlamaConfig.from_pretrained(path)
         config.vocab_size = VOCAB_SIZE
 
         print(f"Loading model from {path}...")
-        self.model = AutoModelForCausalLM.from_pretrained(
+        self.model = LlamaForCausalLM.from_pretrained(
             path,
-            config                = config,
+            config       = config,
             torch_dtype  = torch.float16,
             device_map   = "auto",
-            trust_remote_code = True,
+            # trust_remote_code = True,
             ignore_mismatched_sizes = True, 
         )
         # Resize to match extended vocab (139,336 tokens)