wehub-resource-sync
114 lines
4.8 KiB
Python
114 lines
4.8 KiB
Python
|
|
"""This example shows how to use sentence transformers as a vector embedding model with llmware.
|
|
|
|
To use models from the SentenceTransformer catalog, you may need to install as follows:
|
|
|
|
pip3 install sentence-transformers
|
|
|
|
"""
|
|
|
|
|
|
import os
|
|
|
|
from llmware.setup import Setup
|
|
from llmware.library import Library
|
|
from llmware.retrieval import Query
|
|
from llmware.models import ModelCatalog
|
|
from llmware.configs import LLMWareConfig
|
|
|
|
from importlib import util
|
|
if not util.find_spec("sentence_transformers"):
|
|
print("\nto run this example, you should install the SentenceTransformer library with: "
|
|
"pip3 install sentence-transformers.")
|
|
|
|
def build_lib (library_name, folder="Agreements"):
|
|
|
|
# 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 the documents, text chunks, and captures in MongoDB
|
|
print("update: Step 3 - Parsing and Text Indexing Files")
|
|
|
|
# options: Agreements | UN-Resolutions-500
|
|
library.add_files(input_folder_path=os.path.join(sample_files_path, folder))
|
|
|
|
return library
|
|
|
|
|
|
# start script
|
|
|
|
print("update: Step 1- starting here- building library- parsing PDFs into text chunks")
|
|
|
|
LLMWareConfig().set_active_db("sqlite")
|
|
|
|
lib = build_lib("st_embedding_0")
|
|
|
|
# register a model from the sentence transformers library/repository
|
|
|
|
# note: "all-MiniLM-L6-v2" is from the SentenceTransformer catalog, e.g.,
|
|
# -- https://www.sbert.net/docs/pretrained_models.html
|
|
# -- key inputs to register:
|
|
# -- "model_name" - should be an existing pre-trained model in the SentenceTransformer catalog
|
|
# -- "embedding_dims" - this is the output dimensions, included in the sbert model card info
|
|
# -- "context_window" - included in the sbert model card info
|
|
# -- *** "model_location" - "st_repo" is reserved word to tell llmware to look in sentence transformers ***
|
|
# -- *** "model_family" - "LLMWareSemanticModel" - knows how to load and embed with sentence transformers ***
|
|
|
|
# another sentence transformer to try: "all-mpnet-base-v2" - embedding_dims = 768 - context_window = 384
|
|
|
|
sentence_transformer_pretrained_model_name = "all-MiniLM-L6-v2"
|
|
embedding_dims = 384
|
|
context_window = 256
|
|
|
|
ModelCatalog().register_sentence_transformer_model(model_name=sentence_transformer_pretrained_model_name,
|
|
embedding_dims=embedding_dims, context_window=context_window)
|
|
|
|
"""
|
|
ModelCatalog().add_model_list({"model_name": sentence_transformer_pretrained_model_name,
|
|
"embedding_dims":embedding_dims,
|
|
"context_window":context_window,
|
|
"model_category": "embedding",
|
|
"model_family": "LLMWareSemanticModel",
|
|
"display_name": "MySentenceTransformer", "model_location": "st_repo"})
|
|
"""
|
|
|
|
# to confirm that model has been added to the catalog
|
|
mc = ModelCatalog().list_all_models()
|
|
model_card = ModelCatalog().lookup_model_card(sentence_transformer_pretrained_model_name)
|
|
print("update: model card - ", model_card)
|
|
|
|
# use directly now as an embedding model
|
|
lib.install_new_embedding(embedding_model_name=sentence_transformer_pretrained_model_name,
|
|
vector_db="milvus",batch_size=300)
|
|
|
|
# optional - check the status of the library card and embedding
|
|
lib_card = lib.get_library_card()
|
|
print("update: -- after embedding process - check updated library card - ", lib_card)
|
|
|
|
# create query object (note: including embedding_model is optional - only needed if multiple embeddings on library)
|
|
query_st = Query(lib, embedding_model_name=sentence_transformer_pretrained_model_name)
|
|
|
|
# run multiple queries using query_pgv
|
|
my_search_results = query_st.semantic_query("What is the sale bonus?", result_count = 24)
|
|
|
|
for i, qr in enumerate(my_search_results):
|
|
print("update: semantic query results: ", i, qr)
|
|
|
|
# if you want to delete the embedding - uncomment the line below - including the model_name and vector_db
|
|
# lib.delete_installed_embedding(sentence_transformer_pretrained_model_name, "milvus")
|
|
|
|
# optional - check the embeddings on the library
|
|
emb_record = lib.get_embedding_status()
|
|
for j, entries in enumerate(emb_record):
|
|
print("update: embeddings on library: ", j, entries)
|
|
|