File size: 7,817 Bytes
6330947
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
172
173
174
175
176
177
178
179
180
181
182
183
import logging
import os
from langchain.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import DirectoryLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from dotenv import load_dotenv

from langchain.retrievers import BM25Retriever, EnsembleRetriever
from langchain.llms import OpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain.chains import LLMChain, HypotheticalDocumentEmbedder

## Setting up Log configuration
logging.basicConfig(
    filename='Logs/chatbot.log', # Name of the log file
    level=logging.INFO,      # Logging level (you can use logging.DEBUG for more detailed logs)
    format='%(asctime)s - %(levelname)s - %(message)s'
)


class Jine:

    def __init__(self, OPENAI_API_KEY, VECTOR_STORE_DIRECTORY, VECTOR_STORE_CHECK, DATA_DIRECTORY, DEBUG,USE_HYDE=False):
        self.OPENAI_API_KEY = OPENAI_API_KEY
        self.DATA_DIRECTORY = DATA_DIRECTORY
        self.VECTOR_STORE_DIRECTORY = VECTOR_STORE_DIRECTORY
        self.VECTOR_STORE_CHECK = VECTOR_STORE_CHECK
        self.DEBUG = DEBUG
        self.vectorstore = None
        self.bot = None
        self.USE_HYDE = USE_HYDE
        # creating this variable for BM25 Retriver.
        # self.docs = None

    def create_vectorstore(self):

        if  self.VECTOR_STORE_CHECK:
            print("Loading Vectorstore")
            self.load_vectorstore()
            print('im running')
        else:
            print("Creating Vectorstore")
            docs = DirectoryLoader(self.DATA_DIRECTORY).load()
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=10)
            all_splits = text_splitter.split_documents(docs)
            if self.USE_HYDE:
                base_embeddings = OpenAIEmbeddings()
                llm = OpenAI()
                embeddings_hyde = HypotheticalDocumentEmbedder.from_llm(llm, base_embeddings, "web_search")
                self.vectorstore = Chroma.from_documents(documents=all_splits, embedding=embeddings_hyde,
                                                    persist_directory=self.VECTOR_STORE_DIRECTORY)
                
            else:
                self.vectorstore = Chroma.from_documents(documents=all_splits, embedding=OpenAIEmbeddings(),
                                                    persist_directory=self.VECTOR_STORE_DIRECTORY)


    def multi_query_retriever(self):
        retriever_from_llm = MultiQueryRetriever.from_llm(retriever=self.vectorstore.as_retriever(),
                                                        llm=ChatOpenAI(temperature=0))
        template = """Use the following pieces of context to answer the question at the end. 
        If you don't know the answer, just say that "i am unable to answer your query, for more information contact your HRBP", don't try to make up an answer. 
        Use three sentences maximum and keep the answer as concise as possible.
        {context}
        Question: {question}
        Helpful Answer:"""
        QA_CHAIN_PROMPT = PromptTemplate.from_template(template)

        llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

        self.bot = RetrievalQA.from_chain_type(
            llm,
            retriever=retriever_from_llm,
            chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
        )

    
    def single_query_retriever(self):
        template = """Use the following pieces of context to answer the question at the end. 
        If you don't know the answer, just say that "i am unable to answer your query, for more information contact your HRBP", don't try to make up an answer. 
        Use three sentences maximum and keep the answer as concise as possible.
        {context}
        Question: {question}
        Helpful Answer:"""
        QA_CHAIN_PROMPT = PromptTemplate.from_template(template)

        llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

        self.bot = RetrievalQA.from_chain_type(
        llm,
        retriever=self.vectorstore.as_retriever(),
        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})

    def load_vectorstore(self):
        if self.USE_HYDE:
            print("Using HYDE embeddings vectorstore")
            base_embeddings = OpenAIEmbeddings()
            llm = OpenAI()
            embeddings_hyde = HypotheticalDocumentEmbedder.from_llm(llm, base_embeddings, "web_search")
            self.vectorstore = Chroma(persist_directory=self.VECTOR_STORE_DIRECTORY, embedding_function=embeddings_hyde)
        else:
            print("Using Simple embeddings vectorstore")
            self.vectorstore = Chroma(persist_directory=self.VECTOR_STORE_DIRECTORY, embedding_function=OpenAIEmbeddings())

    def log(self, user_question, chatbot_reply):
        # Log the user's question
        logging.info(f"User: {user_question}")
        # Log the chatbot's reply
        logging.info(f"JIN-e: {chatbot_reply}")

    def load_model(self):
          self.create_vectorstore()
        #   self.multi_query_retriever()
        #   self.single_query_retriever()
          self.create_ensemble_retriever()

    def chat(self, user_question):
        result = self.bot({"query": user_question})
        response = result["result"]
        self.log(user_question, response)
        return response

    ### Adding Ensemble retriver 
    def create_ensemble_retriever(self):
        template = """Use the following pieces of context to answer the question at the end. 
        If you don't know the answer, just say that "i am unable to answer your query, for more information contact your HRBP", don't try to make up an answer. 
        Use three sentences maximum and keep the answer as concise as possible.
        {context}
        Question: {question}
        Helpful Answer:"""

        QA_CHAIN_PROMPT = PromptTemplate.from_template(template)
        print("====================="*10)
        print("Loading Documents for Ensemble Retriver")
        print("====================="*10)

        docs = DirectoryLoader(self.DATA_DIRECTORY).load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=10)
        all_splits = text_splitter.split_documents(docs)

        bm25_retriever = BM25Retriever.from_documents(all_splits)
        # GEttting only two relevant documents 
        bm25_retriever.k = 2

        ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, 
                                                self.vectorstore.as_retriever(search_kwargs={"k": 2})],
                                                weights=[0.5, 0.5])
        
        llm = ChatOpenAI(model_name="gpt-3.5-turbo-16k", temperature=0)

        self.bot = RetrievalQA.from_chain_type(
        llm,
        retriever=ensemble_retriever,
        chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})



if __name__ == "__main__":
    # Set your configuration here
    load_dotenv()
    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
    DATA_DIRECTORY = os.getenv("DATA_DIRECTORY")
    VECTOR_STORE_DIRECTORY = os.getenv("VECTOR_STORE_DIRCTORY")
    VECTOR_STORE_CHECK = os.getenv("VECTOR_STORE_CHECK")
    
    DEBUG = os.getenv("DEBUG")
    USE_HYDE = os.getenv("USE_HYDE")
    # Initialize Jine and start chatting
    jine = Jine(OPENAI_API_KEY, VECTOR_STORE_DIRECTORY, VECTOR_STORE_CHECK, DATA_DIRECTORY, DEBUG)
    # print(jine.VECTOR_STORE_CHECK)
    jine.load_model()
    while True:
        user_question = input("You: ")
        if user_question.lower() in ["exit", "quit"]:
            break
        response = jine.chat(user_question)
        print("JIN-e:", response)