Update for v0.2.2

- Support multiple JS scripts
- Fixed some of bugs
- Resolved a few issue relevant to Colab installation
This commit is contained in:
unclecode
2024-06-02 15:40:18 +08:00
parent f1b60b2016
commit 51f26d12fe
8 changed files with 79 additions and 32 deletions

View File

@@ -188,14 +188,15 @@ class CosineStrategy(ExtractionStrategy):
if self.verbose:
print(f"[LOG] Loading Extraction Model for {self.device.type} device.")
if False and self.device.type == "cpu":
self.model = load_onnx_all_MiniLM_l6_v2()
self.tokenizer = self.model.tokenizer
self.get_embedding_method = "direct"
else:
self.tokenizer, self.model = load_bge_small_en_v1_5()
self.model.eval()
self.get_embedding_method = "batch"
# if False and self.device.type == "cpu":
# self.model = load_onnx_all_MiniLM_l6_v2()
# self.tokenizer = self.model.tokenizer
# self.get_embedding_method = "direct"
# else:
self.tokenizer, self.model = load_bge_small_en_v1_5()
self.model.eval()
self.get_embedding_method = "batch"
self.buffer_embeddings = np.array([])