Update models/qa_model.py

models/qa_model.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from transformers import AutoModelForQuestionAnswering, pipeline
 from features.text_utils import post_process_answer
 from features.graph_utils import find_best_cluster
-
+from optimum.onnxruntime import ORTModelForQuestionAnswering
 
 class QAEnsembleModel(nn.Module):
 
@@ -13,7 +13,7 @@ class QAEnsembleModel(nn.Module):
         super(QAEnsembleModel, self).__init__()
         self.nlps = []
         for model_checkpoint in model_checkpoints:
-            model = AutoModelForQuestionAnswering.from_pretrained(model_name)#.half()
+            model = ORTModelForQuestionAnswering.from_pretrained(model_name, from_transformers=True)#.half()
             model.load_state_dict(torch.load(model_checkpoint, map_location=torch.device('cpu')), strict=False)
             nlp = pipeline('question-answering', model=model,
                            tokenizer=model_name, device=device)