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 AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import logging
 
 logger = logging.getLogger(__name__)
@@ -43,7 +43,7 @@ def model_fn(model_dir, context=None):
         tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
         
         # Load base model
-        base_model = AutoModelForCausalLM.from_pretrained(
+        base_model = AutoModelForSequenceClassification.from_pretrained(
             model_id,
             torch_dtype=torch.bfloat16 if device.type == 'cuda' else torch.float32,
             trust_remote_code=True