|
| 1 | + |
| 2 | +# *** FAST START to create vector embeddings from documents *** |
| 3 | +# |
| 4 | +# docs2vecs_with_milvus-un_resolutions - parses, text chunks and embeds 500 United Nations (UN) Resolutions |
| 5 | +# the sample documents (500 PDFs - 2-15 pages each) can be pulled down from a public S3 repo with the command: |
| 6 | +# sample_files_path = Setup().load_sample_files() |
| 7 | + |
| 8 | +# note: the example assumes that you have installed Milvus and MongoDB per the separate instructions in the README |
| 9 | + |
| 10 | + |
| 11 | +import os |
| 12 | +from llmware.library import Library |
| 13 | +from llmware.retrieval import Query |
| 14 | +from llmware.setup import Setup |
| 15 | +from llmware.status import Status |
| 16 | + |
| 17 | + |
| 18 | +def parse_and_generate_vector_embeddings(library_name): |
| 19 | + |
| 20 | + # Step 0 - Configuration - we will use these in Step 4 to install the embeddings |
| 21 | + embedding_model = "mini-lm-sbert" |
| 22 | + vector_db = "milvus" |
| 23 | + |
| 24 | + # Step 1 - Create library which is the main 'organizing construct' in llmware |
| 25 | + print ("\nupdate: Step 1 - Creating library: {}".format(library_name)) |
| 26 | + |
| 27 | + library = Library().create_new_library(library_name) |
| 28 | + |
| 29 | + # Step 2 - Pull down the sample files from S3 through the .load_sample_files() command |
| 30 | + # --note: if you need to refresh the sample files, set 'over_write=True' |
| 31 | + print ("update: Step 2 - Downloading Sample Files") |
| 32 | + |
| 33 | + sample_files_path = Setup().load_sample_files(over_write=False) |
| 34 | + |
| 35 | + # Step 3 - point ".add_files" method to the folder of documents that was just created |
| 36 | + # this method parses all of the documents, text chunks, and captures in MongoDB |
| 37 | + print("update: Step 3 - Parsing and Text Indexing Files") |
| 38 | + |
| 39 | + library.add_files(input_folder_path=os.path.join(sample_files_path, "UN-Resolutions-500")) |
| 40 | + |
| 41 | + # Step 4 - Install the embeddings |
| 42 | + print("\nupdate: Step 4 - Generating Embeddings in {} db - with Model- {}".format(vector_db, embedding_model)) |
| 43 | + |
| 44 | + library.install_new_embedding(embedding_model_name=embedding_model, vector_db=vector_db) |
| 45 | + |
| 46 | + # note: for using llmware as part of a larger application, you can check the real-time status by polling Status() |
| 47 | + # --both the EmbeddingHandler and Parsers write to Status() at intervals while processing |
| 48 | + update = Status().get_embedding_status(library_name, embedding_model) |
| 49 | + print("update: Embeddings Complete - Status() check at end of embedding - ", update) |
| 50 | + |
| 51 | + # Step 5 - start using the new vector embeddings with Query |
| 52 | + sample_query = "sustainability issues impacting women" |
| 53 | + print("\n\nupdate: Step 5 - Query: {}".format(sample_query)) |
| 54 | + |
| 55 | + query_results = Query(library).semantic_query(sample_query, result_count=20) |
| 56 | + |
| 57 | + for i, entries in enumerate(query_results): |
| 58 | + |
| 59 | + # each query result is a dictionary with many useful keys |
| 60 | + |
| 61 | + text = entries["text"] |
| 62 | + document_source = entries["file_source"] |
| 63 | + page_num = entries["page_num"] |
| 64 | + vector_distance = entries["distance"] |
| 65 | + |
| 66 | + # for display purposes only, we will only show the first 100 characters of the text |
| 67 | + if len(text) > 125: text = text[0:125] + " ... " |
| 68 | + |
| 69 | + print("\nupdate: query results - {} - document - {} - page num - {} - distance - {} " |
| 70 | + .format( i, document_source, page_num, vector_distance)) |
| 71 | + |
| 72 | + print("update: text sample - ", text) |
| 73 | + |
| 74 | + |
| 75 | +if __name__ == "__main__": |
| 76 | + |
| 77 | + # pick any name for the library |
| 78 | + user_selected_name = "un_resolutions500" |
| 79 | + parse_and_generate_vector_embeddings(user_selected_name) |
| 80 | + |
| 81 | + |
0 commit comments