wehub-resource-sync
114 lines
4.0 KiB
Python
114 lines
4.0 KiB
Python
|
|
""" Test for embedding vector creation and storage in a selected vector DB with selected embedding model. """
|
|
|
|
|
|
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 setup_library(library_name):
|
|
|
|
""" Note: this setup_library method is provided to enable a self-contained example to create a test library """
|
|
|
|
# Step 1 - Create library which is the main 'organizing construct' in llmware
|
|
print ("\nupdate: Creating library: {}".format(library_name))
|
|
|
|
library = Library().create_new_library(library_name)
|
|
|
|
# check the embedding status 'before' installing the embedding
|
|
embedding_record = library.get_embedding_status()
|
|
print("embedding record - before embedding ", embedding_record)
|
|
|
|
# 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: 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 the documents, text chunks, and captures in database
|
|
|
|
print("update: Parsing and Text Indexing Files")
|
|
|
|
library.add_files(input_folder_path=os.path.join(sample_files_path, "Agreements"),
|
|
chunk_size=400, max_chunk_size=600, smart_chunking=1)
|
|
|
|
return library
|
|
|
|
|
|
def test_install_vector_embeddings():
|
|
|
|
LLMWareConfig().set_active_db("sqlite")
|
|
|
|
library = setup_library("test_emb_install_09123")
|
|
|
|
# select vector db that you would like to test
|
|
vector_db = "chromadb"
|
|
|
|
LLMWareConfig().set_vector_db(vector_db)
|
|
|
|
# select embedding model
|
|
embedding_model = "mini-lm-sbert"
|
|
|
|
library_name = library.library_name
|
|
|
|
print(f"\nupdate: Starting the Embedding: "
|
|
f"library - {library_name} - "
|
|
f"vector_db - {vector_db} - "
|
|
f"model - {embedding_model}")
|
|
|
|
# *** this is the one key line of code to create the embedding ***
|
|
library.install_new_embedding(embedding_model_name=embedding_model, vector_db=vector_db,batch_size=100)
|
|
|
|
# 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)
|
|
|
|
# Start using the new vector embeddings with Query
|
|
sample_query = "incentive compensation"
|
|
print("\n\nupdate: Run a sample semantic/vector query: {}".format(sample_query))
|
|
|
|
# queries are constructed by creating a Query object, and passing a library as input
|
|
query_results = Query(library).semantic_query(sample_query, result_count=20)
|
|
|
|
assert query_results is not None
|
|
|
|
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"]
|
|
|
|
# to see all of the dictionary keys returned, uncomment the line below
|
|
# print("update: query_results - all - ", i, entries)
|
|
|
|
# for display purposes only, we will only show the first 125 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)
|
|
|
|
# lets take a look at the library embedding status again at the end to confirm embeddings were created
|
|
embedding_record = library.get_embedding_status()
|
|
|
|
assert embedding_record is not None
|
|
|
|
print("\nupdate: embedding record - ", embedding_record)
|
|
|
|
return 0
|
|
|
|
|
|
|
|
|
|
|