|
| 1 | +# Licensed to Elasticsearch B.V. under one or more contributor |
| 2 | +# license agreements. See the NOTICE file distributed with |
| 3 | +# this work for additional information regarding copyright |
| 4 | +# ownership. Elasticsearch B.V. licenses this file to you under |
| 5 | +# the Apache License, Version 2.0 (the "License"); you may |
| 6 | +# not use this file except in compliance with the License. |
| 7 | +# You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | + |
| 18 | +""" |
| 19 | +Vector database index creation example |
| 20 | +""" |
| 21 | + |
| 22 | +import argparse |
| 23 | +import asyncio |
| 24 | +import json |
| 25 | +import os |
| 26 | +from urllib.request import urlopen |
| 27 | + |
| 28 | +import nltk |
| 29 | +from sentence_transformers import SentenceTransformer |
| 30 | +from tqdm import tqdm |
| 31 | + |
| 32 | +from elasticsearch_dsl import ( |
| 33 | + AsyncDocument, |
| 34 | + Date, |
| 35 | + DenseVector, |
| 36 | + InnerDoc, |
| 37 | + Keyword, |
| 38 | + Nested, |
| 39 | + Text, |
| 40 | + async_connections, |
| 41 | +) |
| 42 | + |
| 43 | +DATASET_URL = "https://github.com/raw/elastic/elasticsearch-labs/main/datasets/workplace-documents.json" |
| 44 | +MODEL_NAME = "all-MiniLM-L6-v2" |
| 45 | + |
| 46 | +# initialize sentence tokenizer |
| 47 | +nltk.download("punkt", quiet=True) |
| 48 | + |
| 49 | + |
| 50 | +class Passage(InnerDoc): |
| 51 | + content = Text() |
| 52 | + embedding = DenseVector() |
| 53 | + |
| 54 | + |
| 55 | +class WorkplaceDoc(AsyncDocument): |
| 56 | + class Index: |
| 57 | + name = "workplace_documents" |
| 58 | + |
| 59 | + name = Text() |
| 60 | + summary = Text() |
| 61 | + content = Text() |
| 62 | + created = Date() |
| 63 | + updated = Date() |
| 64 | + url = Keyword() |
| 65 | + category = Keyword() |
| 66 | + passages = Nested(Passage) |
| 67 | + |
| 68 | + _model = None |
| 69 | + |
| 70 | + @classmethod |
| 71 | + def get_embedding_model(cls): |
| 72 | + if cls._model is None: |
| 73 | + cls._model = SentenceTransformer(MODEL_NAME) |
| 74 | + return cls._model |
| 75 | + |
| 76 | + def clean(self): |
| 77 | + # split the content into sentences |
| 78 | + passages = nltk.sent_tokenize(self.content) |
| 79 | + |
| 80 | + # generate an embedding for each passage and save it as a nested document |
| 81 | + model = self.get_embedding_model() |
| 82 | + for passage in passages: |
| 83 | + self.passages.append( |
| 84 | + Passage(content=passage, embedding=list(model.encode(passage))) |
| 85 | + ) |
| 86 | + |
| 87 | + |
| 88 | +async def create(): |
| 89 | + |
| 90 | + # create the index |
| 91 | + await WorkplaceDoc._index.delete(ignore_unavailable=True) |
| 92 | + await WorkplaceDoc.init() |
| 93 | + |
| 94 | + # download the data |
| 95 | + dataset = json.loads(urlopen(DATASET_URL).read()) |
| 96 | + |
| 97 | + # import the dataset |
| 98 | + for data in tqdm(dataset, desc="Indexing documents..."): |
| 99 | + doc = WorkplaceDoc( |
| 100 | + name=data["name"], |
| 101 | + summary=data["summary"], |
| 102 | + content=data["content"], |
| 103 | + created=data.get("created_on"), |
| 104 | + updated=data.get("updated_at"), |
| 105 | + url=data["url"], |
| 106 | + category=data["category"], |
| 107 | + ) |
| 108 | + await doc.save() |
| 109 | + |
| 110 | + |
| 111 | +async def search(query): |
| 112 | + model = WorkplaceDoc.get_embedding_model() |
| 113 | + search = WorkplaceDoc.search().knn( |
| 114 | + field="passages.embedding", |
| 115 | + k=5, |
| 116 | + num_candidates=50, |
| 117 | + query_vector=list(model.encode(query)), |
| 118 | + inner_hits={"size": 3}, |
| 119 | + ) |
| 120 | + async for hit in search: |
| 121 | + print(f"Document: {hit.name}") |
| 122 | + for passage in hit.meta.inner_hits.passages: |
| 123 | + print(f" - [Score: {passage.meta.score}] {passage.content!r}") |
| 124 | + print("") |
| 125 | + |
| 126 | + |
| 127 | +def parse_args(): |
| 128 | + parser = argparse.ArgumentParser(description="Vector database with Elasticsearch") |
| 129 | + parser.add_argument( |
| 130 | + "--create", action="store_true", help="Create and populate a new index" |
| 131 | + ) |
| 132 | + parser.add_argument("query", action="store", help="The search query") |
| 133 | + return parser.parse_args() |
| 134 | + |
| 135 | + |
| 136 | +async def main(): |
| 137 | + args = parse_args() |
| 138 | + |
| 139 | + # initiate the default connection to elasticsearch |
| 140 | + async_connections.create_connection(hosts=[os.environ["ELASTICSEARCH_URL"]]) |
| 141 | + |
| 142 | + if args.create or not await WorkplaceDoc._index.exists(): |
| 143 | + await create() |
| 144 | + |
| 145 | + await search(args.query) |
| 146 | + |
| 147 | + # close the connection |
| 148 | + await async_connections.get_connection().close() |
| 149 | + |
| 150 | + |
| 151 | +if __name__ == "__main__": |
| 152 | + asyncio.run(main()) |
0 commit comments