File size: 7,010 Bytes
097caae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import os
from langchain.document_loaders import PyPDFLoader, DirectoryLoader, PDFMinerLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.vectorstores import Chroma
import configparser
from tqdm import tqdm
from langchain.vectorstores import Pinecone
from langchain.schema import Document
import pinecone
from dotenv import load_dotenv
from llm import LLMManager
class EmbeddingsManager:
def __init__(self,settings, emb="hkunlp/instructor-large"):
#Loading env variables
load_dotenv()
#Loading config file
self.config=configparser.ConfigParser()
self.config.read("config.ini")
#Loading settingManager
self.set=settings
#Loading default parameters for search
self.search_method=self.set.search_method
self.n_doc_return=self.set.n_doc_return
self.ai_assisted_search=self.config.getboolean('RAG','default_ai_assisted_search')
self.available_search_methods=self.set.available_search_methods
self.text_split_size=self.config.getint('RAG','default_text_split_size')
self.text_overlap=self.config.getint('RAG','default_text_overlap')
#Loading available Vector Stores
self.vector_stores=self.get_vector_list()
self.vector_stores_map=self.get_vector_map_list()
#Selecting the embeddings model
self.embedding_model_name=emb
#Initing
current_dir = os.path.dirname(__file__)
data_dir = os.path.join(current_dir, "data")
os.environ['TRANSFORMERS_CACHE'] = data_dir
PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
PINECONE_API_ENV = os.environ.get('PINECONE_API_ENV')
self.embeddings_model = SentenceTransformerEmbeddings(model_name=self.embedding_model_name, cache_folder=data_dir)
pinecone.init(api_key=PINECONE_API_KEY,environment=PINECONE_API_ENV)
#This function used to get the list of emb
def get_emb_list(self):
"""Returns a list of the available Embedding models"""
emb_map_section = 'EMB'
if emb_map_section in self.config:
return [self.config.get(emb_map_section, emb) for emb in self.config[emb_map_section]]
else:
return []
#This function used to get the list of available VectorStores
def get_vector_list(self):
"""Returns a list of the available Vector Stores"""
section = 'Vector_Stores'
if section in self.config:
return [self.config.get(section, vector) for vector in self.config[section]]
else:
return []
#This function used to get the map of available VectorStores
def get_vector_map_list(self):
"""Returns a list of the available Vector Stores"""
section = 'Vector_Stores_Map'
if section in self.config:
return [self.config.get(section, vector) for vector in self.config[section]]
else:
return []
#This function is used to get the relevant context
def get_context(self,index, query, history):
"""Returns the relevant context for the LLM"""
docsearch = Pinecone.from_existing_index(index, self.embeddings_model)
if self.set.ai_assisted_search:
prompt=self.set.default_ai_search_prompt
prompt=prompt.format(question=query,history=history)
print(prompt)
llm=LLMManager(self.set)
queryterms=llm.get_query_terms(prompt)
query=queryterms+"\n"+query
#print("new query input={new_query}".format(new_query=query))
if self.set.search_method=="MMR":
return docsearch.max_marginal_relevance_search(query, k=self.set.n_doc_return,fetch_metadata=True)
elif self.set.search_method=="Similarity":
return docsearch.similarity_search(query, k=self.set.n_doc_return,fetch_metadata=True)
else:
return docsearch.max_marginal_relevance_search(query, k=self.set.n_doc_return,fetch_metadata=True)
#This function is used to get the relevant context
def get_context_search(self,index, query):
"""Returns the relevant context for the LLM"""
docsearch = Pinecone.from_existing_index(index, self.embeddings_model)
if self.set.search_method=="MMR":
return docsearch.max_marginal_relevance_search(query, k=2,fetch_metadata=True)
elif self.set.search_method=="Similarity":
return docsearch.similarity_search(query, k=2,fetch_metadata=True)
else:
return docsearch.max_marginal_relevance_search(query, k=self.n_doc_return,fetch_metadata=True)
#This function is used to get the relevant context formatted
def get_formatted_context(self,index, query,history):
"""Returns the relevant context for the LLM formatted"""
formatted=""
docs=self.get_context(index, query,history)
for doc in docs:
formatted+="DOCUMENT NAME={doc_name}\nDOCUMENT CONTENT={doc_content}\n\n".format(doc_name=doc.metadata["source"],doc_content=doc.page_content)
return formatted
#This function is used to add documents to an existing vector store
def generate_vector_store(self, index):
"""Adds a document to the vector store on Pinecone."""
documents = []
for root, dirs, files in os.walk("docs"):
for file in files:
if file.endswith(".pdf"):
print("Uploading "+file.replace(".pdf",""))
documents.clear()
loader = PDFMinerLoader(os.path.join(root, file))
documents.extend(loader.load())
text_splitter = RecursiveCharacterTextSplitter(chunk_size=self.text_split_size, chunk_overlap=self.text_overlap)
texts = text_splitter.split_documents(documents)
docsearch = Pinecone.from_documents(texts, embedding=self.embeddings_model, index_name=index)
os.remove(os.path.join(root, file))
return "Ok"
# Example Usage:
if __name__ == "__main__":
"""This is an example of how to add document to the vectorstore on Pinecone"""
from settings import SettingManager
set= SettingManager()
emb_manager = EmbeddingsManager(set,emb="hkunlp/instructor-large")
print(emb_manager.generate_vector_store("prohelper"))
#"""This is an example of how to retrive context and display all values retrived"""
#emb_manager = EmbeddingsManager()
#docs=emb_manager.get_context(index="prohelper",query="Could you explain to me what is esrs?")
#for i in docs:
# print("---------------------------------------------------------")
# print(i.metadata["Doc"])
# print(" ")
# print(i.page_content)
#llm_manager.selectLLM("Mixtral 7B") |