Spaces:
Running
Running
| """ | |
| /************************************************************************* | |
| * | |
| * CONFIDENTIAL | |
| * __________________ | |
| * | |
| * Copyright (2023-2024) AI Labs, IronOne Technologies, LLC | |
| * All Rights Reserved | |
| * | |
| * Author : Theekshana Samaradiwakara | |
| * Description :Python Backend API to chat with private data | |
| * CreatedDate : 14/11/2023 | |
| * LastModifiedDate : 18/03/2024 | |
| *************************************************************************/ | |
| """ | |
| import os | |
| import time | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| from dotenv import load_dotenv | |
| from reggpt.chains.llmChain import get_qa_chain | |
| from reggpt.output_parsers.output_parser import qa_chain_output_parser | |
| from reggpt.configs.model import QA_MODEL_TYPE | |
| from reggpt.utils.retriever import load_ensemble_retriever | |
| load_dotenv() | |
| verbose = os.environ.get('VERBOSE') | |
| qa_model_type=QA_MODEL_TYPE | |
| # retriever=load_faiss_retriever() | |
| retriever=load_ensemble_retriever() | |
| # retriever=load_multi_query_retriever(multi_query_model_type) | |
| logger.info("retriever loaded:") | |
| qa_chain= get_qa_chain(qa_model_type,retriever) | |
| def run_qa_chain(query): | |
| try: | |
| logger.info(f"run_qa_chain : Question: {query}") | |
| # Get the answer from the chain | |
| start = time.time() | |
| # res = qa_chain(query) | |
| res = qa_chain.invoke({"question": query, "chat_history":""}) | |
| if 'I dont know' in res["answer"] or "don't know" in res["answer"]: | |
| res['answer'] = "I currently do not have the information to answer this question. Please rephrase your question or ask another question." | |
| # res = response | |
| # answer, docs = res['result'],res['source_documents'] | |
| end = time.time() | |
| # log the result | |
| logger.info(f"Answer (took {round(end - start, 2)} s.) \n: {res}") | |
| return qa_chain_output_parser(res) | |
| except Exception as e: | |
| logger.exception(e) | |
| raise e | |
| def run_qa_chain_answer_only(query): | |
| try: | |
| logger.info(f"run_qa_chain : Question: {query}") | |
| # Get the answer from the chain | |
| start = time.time() | |
| # res = qa_chain(query) | |
| res = qa_chain.invoke({"question": query, "chat_history":""}) | |
| # res = response | |
| # answer, docs = res['result'],res['source_documents'] | |
| end = time.time() | |
| # log the result | |
| logger.info(f"Answer (took {round(end - start, 2)} s.) \n: {res}") | |
| return res["answer"] | |
| except Exception as e: | |
| logger.exception(e) | |
| raise e |