| 59 | |
| 60 | @singleton |
| 61 | class EmbeddingHelper: |
| 62 | def __init__(self, embeddingPath='./embedding', embeddingFile='test', need_redis=False, embedding_type='default'): |
| 63 | super().__init__() |
| 64 | self.embeddingPath = embeddingPath |
| 65 | self.embeddingFile = embeddingFile |
| 66 | self.client = chromadb.PersistentClient(path=embeddingPath) |
| 67 | # self.client = chromadb.Client() |
| 68 | |
| 69 | self.collections = self.client.list_collections() |
| 70 | |
| 71 | # sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction( |
| 72 | # model_name="all-MiniLM-L6-v2") |
| 73 | default_ef = embedding_functions.DefaultEmbeddingFunction() |
| 74 | |
| 75 | self.embedding_type = embedding_type |
| 76 | |
| 77 | self.embedding_model = default_ef |
| 78 | self.collection = self.client.get_or_create_collection(name=embeddingFile, |
| 79 | embedding_function=default_ef) |
| 80 | |
| 81 | self.image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp'] |
| 82 | |
| 83 | self.need_redis = need_redis |
| 84 | |
| 85 | self.checkClient() |
| 86 | |
| 87 | if self.need_redis: |
| 88 | self.redisCaching = RedisCaching() |
| 89 | |
| 90 | def query(self, text, topK=5, where=None, where_document=None): |
| 91 | # where={"metadata_field": "is_equal_to_this"}, |
| 92 | # where_document={"$contains": "search_string"} |
| 93 | if self.need_redis: |
| 94 | starttime = time.time() |
| 95 | res = self.query_preprocessing(text, top_k=topK) |
| 96 | if res is not None: |
| 97 | if len(res) >= 1: |
| 98 | return res |
| 99 | if where is None: |
| 100 | if where_document is None: |
| 101 | result = self.collection.query( |
| 102 | query_texts=[text], |
| 103 | n_results=topK |
| 104 | ) |
| 105 | else: |
| 106 | result = self.collection.query( |
| 107 | query_texts=[text], |
| 108 | n_results=topK, |
| 109 | ) |
| 110 | else: |
| 111 | if where_document is None: |
| 112 | result = self.collection.query( |
| 113 | query_texts=[text], |
| 114 | n_results=topK, |
| 115 | where=where |
| 116 | ) |
| 117 | else: |
| 118 | result = self.collection.query( |