orionweller commited on
Commit
1239c17
·
1 Parent(s): dc7edc7
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -37,6 +37,7 @@ qrels = {}
37
  datasets = ["scifact"]
38
  current_dataset = "scifact"
39
 
 
40
  def pool(last_hidden_states, attention_mask):
41
  last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
42
  sequence_lengths = attention_mask.sum(dim=1) - 1
@@ -69,8 +70,8 @@ def load_model():
69
  tokenizer.padding_side = "right"
70
 
71
  # model = AutoModel.from_pretrained(CUR_MODEL, max_memory={"cpu": "12GiB"}, torch_dtype=torch.bfloat16, offload_state_dict=True)
72
- base_model_instance = AutoModel.from_pretrained(BASE_MODEL)
73
- model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL)
74
  model.eval()
75
 
76
 
 
37
  datasets = ["scifact"]
38
  current_dataset = "scifact"
39
 
40
+
41
  def pool(last_hidden_states, attention_mask):
42
  last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
43
  sequence_lengths = attention_mask.sum(dim=1) - 1
 
70
  tokenizer.padding_side = "right"
71
 
72
  # model = AutoModel.from_pretrained(CUR_MODEL, max_memory={"cpu": "12GiB"}, torch_dtype=torch.bfloat16, offload_state_dict=True)
73
+ base_model_instance = AutoModel.from_pretrained(BASE_MODEL, device="cpu")
74
+ model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL, device="cpu")
75
  model.eval()
76
 
77