Update src/models/qa_model.py

src/models/qa_model.py

@@ -9,14 +9,14 @@ from src.features.graph_utils import find_best_cluster
 class QAEnsembleModel(nn.Module):
 
     def __init__(self, model_name, model_checkpoints, entity_dict,
-                 thr=0.1, device="CPU:0"):
+                 thr=0.1, device="cpu"):
         super(QAEnsembleModel, self).__init__()
         self.nlps = []
         for model_checkpoint in model_checkpoints:
-            model = AutoModelForQuestionAnswering.from_pretrained(model_name).half()
+            model = AutoModelForQuestionAnswering.from_pretrained(model_name)#.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=int(device.split(":")[-1]))
+                           tokenizer=model_name, device=0)
             self.nlps.append(nlp)
         self.entity_dict = entity_dict
         self.thr = thr