cpu

app.py

@@ -37,6 +37,7 @@ qrels = {}
 datasets = ["scifact"]
 current_dataset = "scifact"
 
+
 def pool(last_hidden_states, attention_mask):
     last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
     sequence_lengths = attention_mask.sum(dim=1) - 1
@@ -69,8 +70,8 @@ def load_model():
     tokenizer.padding_side = "right"
     
     # model = AutoModel.from_pretrained(CUR_MODEL, max_memory={"cpu": "12GiB"}, torch_dtype=torch.bfloat16, offload_state_dict=True)
-    base_model_instance = AutoModel.from_pretrained(BASE_MODEL)
-    model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL)
+    base_model_instance = AutoModel.from_pretrained(BASE_MODEL, device="cpu")
+    model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL, device="cpu")
     model.eval()