Spaces:
Runtime error
Runtime error
| """ | |
| Python Backend API to chat with private data | |
| 08/14/2023 | |
| D.M. Theekshana Samaradiwakara | |
| """ | |
| import os | |
| import time | |
| from dotenv import load_dotenv | |
| from langchain.chains import RetrievalQA | |
| from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler | |
| from langchain.llms import GPT4All | |
| from langchain.llms import HuggingFaceHub | |
| from langchain.chat_models import ChatOpenAI | |
| # from langchain.retrievers.self_query.base import SelfQueryRetriever | |
| # from langchain.chains.query_constructor.base import AttributeInfo | |
| # from chromaDb import load_store | |
| from faissDb import load_FAISS_store | |
| load_dotenv() | |
| #gpt4 all model | |
| gpt4all_model_path = os.environ.get('GPT4ALL_MODEL_PATH') | |
| model_n_ctx = os.environ.get('MODEL_N_CTX') | |
| model_n_batch = int(os.environ.get('MODEL_N_BATCH',8)) | |
| target_source_chunks = int(os.environ.get('TARGET_SOURCE_CHUNKS',4)) | |
| openai_api_key = os.environ.get('OPENAI_API_KEY') | |
| verbose = os.environ.get('VERBOSE') | |
| # activate/deactivate the streaming StdOut callback for LLMs | |
| callbacks = [StreamingStdOutCallbackHandler()] | |
| class QAPipeline: | |
| def __init__(self): | |
| self.llm_name = None | |
| self.llm = None | |
| self.dataset_name = None | |
| self.vectorstore = None | |
| self.qa_chain = None | |
| def run(self,query, model, dataset): | |
| if (self.llm_name != model) or (self.dataset_name != dataset) or (self.qa_chain == None): | |
| self.set_model(model) | |
| self.set_vectorstore(dataset) | |
| self.set_qa_chain() | |
| # Get the answer from the chain | |
| start = time.time() | |
| res = self.qa_chain(query) | |
| # answer, docs = res['result'],res['source_documents'] | |
| end = time.time() | |
| # Print the result | |
| print("\n\n> Question:") | |
| print(query) | |
| print(f"\n> Answer (took {round(end - start, 2)} s.):") | |
| print( res) | |
| return res | |
| def set_model(self,model_type): | |
| if model_type != self.llm_name: | |
| match model_type: | |
| case "gpt4all": | |
| # self.llm = GPT4All(model=gpt4all_model_path, n_ctx=model_n_ctx, backend='gptj', n_batch=model_n_batch, callbacks=callbacks, verbose=verbose) | |
| self.llm = GPT4All(model=gpt4all_model_path, max_tokens=model_n_ctx, backend='gptj', n_batch=model_n_batch, callbacks=callbacks, verbose=verbose) | |
| # self.llm = HuggingFaceHub(repo_id="nomic-ai/gpt4all-j", model_kwargs={"temperature":0.001, "max_length":1024}) | |
| case "google/flan-t5-xxl": | |
| self.llm = HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature":0.001, "max_length":1024}) | |
| case "tiiuae/falcon-7b-instruct": | |
| self.llm = HuggingFaceHub(repo_id=model_type, model_kwargs={"temperature":0.001, "max_length":1024}) | |
| case "openai": | |
| self.llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0) | |
| case _default: | |
| # raise exception if model_type is not supported | |
| raise Exception(f"Model type {model_type} is not supported. Please choose a valid one") | |
| self.llm_name = model_type | |
| def set_vectorstore(self, dataset): | |
| if dataset != self.dataset_name: | |
| # self.vectorstore = load_store(dataset) | |
| self.vectorstore = load_FAISS_store() | |
| print("\n\n> vectorstore loaded:") | |
| self.dataset_name = dataset | |
| def set_qa_chain(self): | |
| self.qa_chain = RetrievalQA.from_chain_type( | |
| llm=self.llm, | |
| chain_type="stuff", | |
| retriever = self.vectorstore.as_retriever(), | |
| # retriever = self.vectorstore.as_retriever(search_kwargs={"k": target_source_chunks} | |
| return_source_documents= True | |
| ) | |