wehub-resource-sync
100 lines
3.8 KiB
Python
100 lines
3.8 KiB
Python
|
|
""" This example illustrates parsing, text chunking, embedding and then executing in a RAG prompt process using
|
|
~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.
|
|
|
|
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.prompts import Prompt
|
|
from llmware.configs import LLMWareConfig
|
|
|
|
|
|
def rag (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)
|
|
contracts_path = os.path.join(sample_files_path, "Agreements")
|
|
|
|
# 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=contracts_path)
|
|
|
|
# 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)
|
|
|
|
print("\nupdate: Loading 1B parameter BLING model for LLM inference")
|
|
|
|
prompter = Prompt().load_model("llmware/bling-1b-0.1")
|
|
query = "what is the executive's base annual salary"
|
|
|
|
results = Query(library).semantic_query(query, result_count=50, embedding_distance_threshold=1.0)
|
|
|
|
for i, res in enumerate(results):
|
|
|
|
print("update: ", i, res["file_source"], res["distance"], res["text"])
|
|
|
|
for i, contract in enumerate(os.listdir(contracts_path)):
|
|
|
|
qr = []
|
|
|
|
if contract != ".DS_Store":
|
|
|
|
print("\nContract Name: ", i, contract)
|
|
|
|
for j, entries in enumerate(results):
|
|
if entries["file_source"] == contract:
|
|
print("Top Retrieval: ", j, entries["distance"], entries["text"])
|
|
qr.append(entries)
|
|
|
|
source = prompter.add_source_query_results(query_results=qr)
|
|
response = prompter.prompt_with_source(query, prompt_name="default_with_context", temperature=0.3)
|
|
|
|
for resp in response:
|
|
if "llm_response" in resp:
|
|
print("\nupdate: llm answer - ", resp["llm_response"])
|
|
|
|
# start fresh for next document
|
|
prompter.clear_source_materials()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
# feel free to change to sqlite or postgres (if installed)
|
|
LLMWareConfig().set_active_db("mongo")
|
|
|
|
# pick any name for the library
|
|
user_selected_name = "contracts_rag10"
|
|
rag(user_selected_name)
|
|
|
|
|