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 AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
+from transformers import AutoModel, AutoTokenizer, AutoConfig
 import logging
 
 logger = logging.getLogger(__name__)
@@ -42,15 +42,13 @@ def model_fn(model_dir, context=None):
         # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
         
-        # Configure model as RoBERTa
+        # Load config and specify it's a Phi3Config
         config = AutoConfig.from_pretrained(model_id, 
                                           num_labels=2,
-                                          architectures=["RobertaForSequenceClassification"],
-                                          model_type="roberta",
                                           trust_remote_code=True)
         
-        # Load base model with RoBERTa config
-        base_model = AutoModelForSequenceClassification.from_pretrained(
+        # Load base model
+        base_model = AutoModel.from_pretrained(
             model_id,
             config=config,
             torch_dtype=torch.bfloat16 if device.type == 'cuda' else torch.float32,