Hackoor/SampleLlamaModel-2-FINAL

1import streamlit as st2from streamlit_chat import message3from langchain.chains import ConversationalRetrievalChain4from langchain.embeddings import HuggingFaceEmbeddings5from langchain.llms import CTransformers6from langchain.llms import Replicate7from langchain.text_splitter import CharacterTextSplitter8from langchain.vectorstores import FAISS9from langchain.memory import ConversationBufferMemory10from langchain.document_loaders import PyPDFLoader11from langchain.document_loaders import TextLoader12from langchain.document_loaders import Docx2txtLoader13from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler14import os15from dotenv import load_dotenv16import tempfile17 18 19load_dotenv()20 21 22def initialize_session_state():23    if 'history' not in st.session_state:24        st.session_state['history'] = []25 26    if 'generated' not in st.session_state:27        st.session_state['generated'] = ["Hello! Ask me anything about 🤗"]28 29    if 'past' not in st.session_state:30        st.session_state['past'] = ["Hey! 👋"]31 32def conversation_chat(query, chain, history):33    result = chain({"question": query, "chat_history": history})34    history.append((query, result["answer"]))35    return result["answer"]36 37def display_chat_history(chain):38    reply_container = st.container()39    container = st.container()40 41    with container:42        with st.form(key='my_form', clear_on_submit=True):43            user_input = st.text_input("Question:", placeholder="Ask about your Documents", key='input')44            submit_button = st.form_submit_button(label='Send')45 46        if submit_button and user_input:47            with st.spinner('Generating response...'):48                output = conversation_chat(user_input, chain, st.session_state['history'])49 50            st.session_state['past'].append(user_input)51            st.session_state['generated'].append(output)52 53    if st.session_state['generated']:54        with reply_container:55            for i in range(len(st.session_state['generated'])):56                message(st.session_state["past"][i], is_user=True, key=str(i) + '_user', avatar_style="thumbs")57                message(st.session_state["generated"][i], key=str(i), avatar_style="fun-emoji")58 59def create_conversational_chain(vector_store):60    load_dotenv()61    # Create llm62    #llm = CTransformers(model="llama-2-7b-chat.ggmlv3.q4_0.bin",63                        #streaming=True, 64                        #callbacks=[StreamingStdOutCallbackHandler()],65                        #model_type="llama", config={'max_new_tokens': 500, 'temperature': 0.01})66    llm = Replicate(67        streaming = True,68        model = "replicate/llama-2-70b-chat:58d078176e02c219e11eb4da5a02a7830a283b14cf8f94537af893ccff5ee781", 69        callbacks=[StreamingStdOutCallbackHandler()],70        input = {"temperature": 0.01, "max_length" :500,"top_p":1})71    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)72 73    chain = ConversationalRetrievalChain.from_llm(llm=llm, chain_type='stuff',74                                                 retriever=vector_store.as_retriever(search_kwargs={"k": 2}),75                                                 memory=memory)76    return chain77 78def main():79    load_dotenv()80    # Initialize session state81    initialize_session_state()82    st.title("Multi-Docs ChatBot using llama-2-70b :books:")83    # Initialize Streamlit84    st.sidebar.title("Document Processing")85    uploaded_files = st.sidebar.file_uploader("Upload files", accept_multiple_files=True)86 87 88    if uploaded_files:89        text = []90        for file in uploaded_files:91            file_extension = os.path.splitext(file.name)[1]92            with tempfile.NamedTemporaryFile(delete=False) as temp_file:93                temp_file.write(file.read())94                temp_file_path = temp_file.name95 96            loader = None97            if file_extension == ".pdf":98                loader = PyPDFLoader(temp_file_path)99            elif file_extension == ".docx" or file_extension == ".doc":100                loader = Docx2txtLoader(temp_file_path)101            elif file_extension == ".txt":102                loader = TextLoader(temp_file_path)103 104            if loader:105                text.extend(loader.load())106                os.remove(temp_file_path)107 108        text_splitter = CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=100, length_function=len)109        text_chunks = text_splitter.split_documents(text)110 111        # Create embeddings112        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", 113                                           model_kwargs={'device': 'cpu'})114 115        # Create vector store116        vector_store = FAISS.from_documents(text_chunks, embedding=embeddings)117 118        # Create the chain object119        chain = create_conversational_chain(vector_store)120 121        122        display_chat_history(chain)123 124if __name__ == "__main__":125    main()126 127