Update models/pairwise_model.py

models/pairwise_model.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from torch.utils.data import Dataset, DataLoader
 from transformers import AutoModel, AutoConfig
 from transformers import AutoTokenizer
-# from optimum.onnxruntime import ORTModelForQuestionAnswering
+from optimum.onnxruntime import ORTModel
 import pandas as pd
 import os
 
@@ -21,7 +21,7 @@ class PairwiseModel(nn.Module):
         self.max_length = max_length
         self.batch_size = batch_size
         self.device = device
-        self.model = AutoModel.from_pretrained(model_name, use_auth_token=AUTH_TOKEN)
+        self.model = ORTModel.from_pretrained(model_name, use_auth_token=AUTH_TOKEN, from_transformers=True)
         self.model.to(self.device)
         self.model.eval()
         self.config = AutoConfig.from_pretrained(model_name, use_auth_token=AUTH_TOKEN)