wehub-resource-sync
85 lines
3.4 KiB
Python
85 lines
3.4 KiB
Python
|
|
""" This example illustrates parsing, text chunking, embedding and then querying ~80 legal documents. The
|
|
example was originally developed for a joint webinar hosted with Milvus. Please feel free to substitute
|
|
other vector databases in the example, if you prefer, along with changing the text collection DB from Mongo to
|
|
either SQLite or Postgres.
|
|
|
|
The example uses sample documents (~80 legal template contracts) that can be pulled down with the command:
|
|
sample_files_path = Setup().load_sample_files()
|
|
"""
|
|
|
|
|
|
import os
|
|
from llmware.library import Library
|
|
from llmware.retrieval import Query
|
|
from llmware.setup import Setup
|
|
from llmware.status import Status
|
|
from llmware.configs import LLMWareConfig
|
|
|
|
|
|
def parse_and_generate_vector_embeddings(library_name):
|
|
|
|
# Step 0 - Configuration - we will use these in Step 4 to install the embeddings
|
|
embedding_model = "industry-bert-contracts"
|
|
vector_db = "milvus"
|
|
|
|
# Step 1 - Create library which is the main 'organizing construct' in llmware
|
|
print ("\nupdate: Step 1 - Creating library: {}".format(library_name))
|
|
|
|
library = Library().create_new_library(library_name)
|
|
|
|
# Step 2 - Pull down the sample files from S3 through the .load_sample_files() command
|
|
# --note: if you need to refresh the sample files, set 'over_write=True'
|
|
print ("update: Step 2 - Downloading Sample Files")
|
|
|
|
sample_files_path = Setup().load_sample_files(over_write=False)
|
|
|
|
# Step 3 - point ".add_files" method to the folder of documents that was just created
|
|
# this method parses all of the documents, text chunks, and captures in MongoDB
|
|
print("update: Step 3 - Parsing and Text Indexing Files")
|
|
|
|
library.add_files(input_folder_path=os.path.join(sample_files_path, "AgreementsLarge"))
|
|
|
|
# Step 4 - Install the embeddings
|
|
print("\nupdate: Step 4 - Generating Embeddings in {} db - with Model- {}".format(vector_db, embedding_model))
|
|
|
|
library.install_new_embedding(embedding_model_name=embedding_model, vector_db=vector_db)
|
|
|
|
# note: for using llmware as part of a larger application, you can check the real-time status by polling Status()
|
|
# --both the EmbeddingHandler and Parsers write to Status() at intervals while processing
|
|
update = Status().get_embedding_status(library_name, embedding_model)
|
|
print("update: Embeddings Complete - Status() check at end of embedding - ", update)
|
|
|
|
# Step 5 - start using the new vector embeddings with Query
|
|
sample_query = "incentive compensation"
|
|
print("\n\nupdate: Step 5 - Query: {}".format(sample_query))
|
|
|
|
query_results = Query(library).semantic_query(sample_query, result_count=20)
|
|
|
|
for i, entries in enumerate(query_results):
|
|
|
|
# each query result is a dictionary with many useful keys
|
|
|
|
text = entries["text"]
|
|
document_source = entries["file_source"]
|
|
page_num = entries["page_num"]
|
|
vector_distance = entries["distance"]
|
|
|
|
# for display purposes only, we will only show the first 100 characters of the text
|
|
if len(text) > 125: text = text[0:125] + " ... "
|
|
|
|
print("\nupdate: query results - {} - document - {} - page num - {} - distance - {} "
|
|
.format( i, document_source, page_num, vector_distance))
|
|
|
|
print("update: text sample - ", text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
LLMWareConfig().set_active_db("mongo")
|
|
|
|
# pick any name for the library
|
|
user_selected_name = "contracts"
|
|
parse_and_generate_vector_embeddings(user_selected_name)
|
|
|