add cash biencode
Browse files- bi_encoder.py +2 -2
bi_encoder.py
CHANGED
@@ -4,8 +4,8 @@ import torch
|
|
4 |
from transformers import AutoTokenizer, AutoModel
|
5 |
from sklearn.metrics.pairwise import cosine_similarity
|
6 |
|
7 |
-
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased", cache_dir="model/")
|
8 |
-
bert_model = AutoModel.from_pretrained("distilbert-base-uncased", cache_dir="model/")
|
9 |
|
10 |
|
11 |
def mean_pool(token_embeds: torch.tensor, attention_mask: torch.tensor) -> torch.tensor:
|
|
|
4 |
from transformers import AutoTokenizer, AutoModel
|
5 |
from sklearn.metrics.pairwise import cosine_similarity
|
6 |
|
7 |
+
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased", cache_dir="/model/transformers_cache")
|
8 |
+
bert_model = AutoModel.from_pretrained("distilbert-base-uncased", cache_dir="/model/transformers_cache")
|
9 |
|
10 |
|
11 |
def mean_pool(token_embeds: torch.tensor, attention_mask: torch.tensor) -> torch.tensor:
|