Update models/pairwise_model.py

models/pairwise_model.py

@@ -4,9 +4,11 @@ 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
 import pandas as pd
+import os
 
-AUTH_TOKEN = "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"
+AUTH_TOKEN = os.getenv("AUTH_TOKEN")
 
 tokenizer = AutoTokenizer.from_pretrained('nguyenvulebinh/vi-mrc-base',
                                           use_auth_token=AUTH_TOKEN)
@@ -19,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 = ORTModelForQuestionAnswering.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)