DishaLLM/QAChatBot

1 2import os3from langchain.text_splitter import CharacterTextSplitter4from langchain.document_loaders import TextLoader, DirectoryLoader5from langchain.embeddings import CohereEmbeddings6from langchain.embeddings import OpenAIEmbeddings7from langchain.vectorstores import Chroma8from langchain.llms import OpenAI9from langchain.llms import Cohere10from langchain.chains import RetrievalQA11from langchain.prompts import PromptTemplate12 13import streamlit as st14 15def retrieve(query,llm,retriever):16 17    template = """18    Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.Use only the document for your answer and you may summarize the answer in 50 words to make it look better.19 20    {context}21 22    Question: {question}23    """24    # create the chain to answer questions25    qa_chain = RetrievalQA.from_chain_type(llm=llm,26                                  chain_type="stuff",27                                  retriever=retriever,28                                  chain_type_kwargs={29                                  "prompt": PromptTemplate(30                                  template=template,31                                  input_variables=["context", "question"],)})32 33    return(qa_chain.run(query))34 35 36def main():37 38    # Main title of the application39    st.title("Q&A BOT")40 41    if 'counter' not in st.session_state:42            st.session_state['counter'] = 043 44    with st.sidebar:45        with st.form('Cohere/OpenAI'):46            mod = st.radio('Choose OpenAI/Cohere', ('OpenAI', 'Cohere'))47            api_key = st.text_input('Enter API key', type="password")48            # model = st.radio('Choose Company', ('ArtisanAppetite foods', 'BMW','Titan Watches'))49            submitted = st.form_submit_button("Submit")50 51    if api_key:52        if(mod=='OpenAI'):53            os.environ["OPENAI_API_KEY"] = api_key54            llm = OpenAI(temperature=0.7, verbose=True)55            embeddings = OpenAIEmbeddings()56        elif(mod=='Cohere'):57            os.environ["COHERE_API_KEY"] = api_key58            llm = Cohere(temperature=0.7, verbose=True)59            embeddings = CohereEmbeddings()60 61        uploaded_file = st.file_uploader("Upload a file to ingest", type=["txt"])62 63        if uploaded_file is not None:64 65 66            file_path = uploaded_file.name67            print(file_path)68 69            # this is a necessary step to read the file content and save it 70            # in the webservers location71            file_contents = uploaded_file.read()72            save_path = uploaded_file.name73            with open(save_path, "wb") as f:74                f.write(file_contents)75            print(save_path)76 77            loader = TextLoader(save_path,autodetect_encoding=True)78            documents = loader.load()79            text_splitter = CharacterTextSplitter(chunk_size=1000) #Splitting the text and creating chunks80            docs = text_splitter.split_documents(documents)81 82            persist_directory = save_path[:-4]83            vectordb = Chroma.from_documents(documents=docs,84                                        embedding=embeddings,85                                        persist_directory=persist_directory)86            # persiste the db to disk87            vectordb.persist()88            retriever = vectordb.as_retriever(search_kwargs={"k": 3})89 90            st.session_state['counter'] += 191 92        if st.session_state['counter'] > 0:93            query = st.text_input("Query: ", "", key="input")94            result_display = st.empty()95 96            if query is not None and query != "":97 98                # create a retriever99                100                result = retrieve(query,llm,retriever)101                # Text area for editing the generated response102                result_display.text_area("Result:", value=result, height=500)103 104    elif (not api_key):105        st.info("Please add configuration details in left panel")106        st.stop() 107 108if __name__ == "__main__":109    main()110 111