Upload chatbot.py

chatbot.py

@@ -0,0 +1,108 @@
+import os
+from langchain_community.chat_models import ChatOpenAI
+
+from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.embeddings import OpenAIEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain.chains import ConversationalRetrievalChain
+
+import streamlit as st
+from streamlit_chat import message
+
+
+@st.cache_data()
+def load_docs():
+    documents = []
+    for file in os.listdir('docs'):
+        if file.endswith('.pdf'):
+            pdf_path = "./docs/"+file
+            loader = PyPDFLoader(pdf_path)
+            documents.extend(loader.load())
+        elif file.endswith('.docx') or file.endswith('.doc'):
+            doc_path = './docs/'+file
+            loader = Docx2txtLoader(doc_path)
+            documents.extend(loader.load())
+        elif file.endswith('.txt'):
+            text_path = '.docs/'+file
+            loader = TextLoader(text_path)
+            documents.extend(loader.load())
+    
+    return documents
+
+os.environ["OPENAI_API_KEY"] = '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'
+
+llm_model = "gpt-3.5-turbo"
+llm = ChatOpenAI(temperature=.7, model=llm_model)
+#======================================================================================================================
+# Load documents
+documents = load_docs()
+chat_history = []
+
+# 1. Text splitter
+text_splitter = CharacterTextSplitter(
+    chunk_size = 100,
+    chunk_overlap = 20,
+    length_function = len
+)
+
+# 2. Embedding
+embeddings = OpenAIEmbeddings()
+
+docs = text_splitter.split_documents(documents)
+
+#=====================================================================================================================
+# 3. Storage
+vector_store = Chroma.from_documents(
+    documents=docs,
+    embedding=embeddings,
+    persist_directory='./data'
+)
+vector_store.persist()
+# ====================================================================================================================
+# 4. Retrieve
+retriever = vector_store.as_retriever(search_kwargs={"k":6})
+# docs = retriever.get_relevant_documents("Tell me more about Data Science")
+
+# Make a chain to answer questions
+qa_chain = ConversationalRetrievalChain.from_llm(
+    llm, 
+    vector_store.as_retriever(search_kwargs={'k':6}),
+    return_source_documents=True,
+    verbose=False
+)
+ 
+
+# cite sources - helper function to prettyfy responses
+def process_llm_response(llm_response):
+    print(llm_response['result'])
+    print('\n\nSources:')
+    for source in llm_response['source_documents']:
+        print(source.metadata['source'])
+        
+#==============================FRONTEND=======================================
+st.title("ViTo chatbot👠")
+st.header("Ask anything about ViTo company...")
+
+if 'generated' not in st.session_state:
+    st.session_state['generated'] = []
+
+if 'past' not in st.session_state:
+    st.session_state['past'] = []
+    
+def get_query():
+    input_text = st.chat_input("Ask a question about your documents...")
+    return input_text
+
+# retrieve the user input
+user_input = get_query()
+if user_input:
+    result = qa_chain({'question': user_input, 'chat_history': chat_history})
+    st.session_state.past.append(user_input)
+    st.session_state.generated.append(result['answer'])
+    
+if st.session_state['generated']:
+     for i in range(len(st.session_state['generated'])):
+         message(st.session_state['past'][i], is_user=True, key=str(i)+'_user')
+         message(st.session_state['generated'][i], key=str(i))
+