Update code/inference.py

code/inference.py

@@ -2,7 +2,7 @@ import os
 import json
 import torch
 import torch.nn as nn
-from transformers import AutoModel, AutoTokenizer, AutoConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import logging
 
 logger = logging.getLogger(__name__)
@@ -42,15 +42,9 @@ def model_fn(model_dir, context=None):
         # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
         
-        # Load config and specify it's a Phi3Config
-        config = AutoConfig.from_pretrained(model_id, 
-                                          num_labels=2,
-                                          trust_remote_code=True)
-        
         # Load base model
-        base_model = AutoModel.from_pretrained(
+        base_model = AutoModelForCausalLM.from_pretrained(
             model_id,
-            config=config,
             torch_dtype=torch.bfloat16 if device.type == 'cuda' else torch.float32,
             trust_remote_code=True
         )