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