Spaces:
Paused
Paused
| import streamlit as st | |
| from run_llama import load_models | |
| import variables as vr | |
| from load_documents import load_documents_fn | |
| import torch | |
| from langchain.vectorstores import Chroma | |
| from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings | |
| from langchain.chains import RetrievalQA | |
| from langchain.prompts import PromptTemplate | |
| from langchain.memory import ConversationBufferMemory | |
| from variables import EMBEDDING_MODEL_NAME,MODEL_ID,MODEL_BASENAME | |
| from chromadb.utils import embedding_functions | |
| # print(f"Is CUDA available: {torch.cuda.is_available()}") | |
| # # True | |
| # print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}") | |
| # Tesla T4 | |
| def model_memory(): | |
| # Adding history to the model. | |
| template = """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. | |
| {context} | |
| Question: {question} | |
| Helpful Answer:""" | |
| # template = """Use the following pieces of context to answer the question at the end. If you don't know the answer,\ | |
| # try to make up an short and sweet answer. | |
| # if the context does not exists or has no relevancy then try answer it on your previous knowledge base.But keep the answer short and precise. | |
| # {context} | |
| # | |
| # | |
| # Question: {question} | |
| # Helpful Answer:""" | |
| prompt = PromptTemplate(input_variables=["context", "question"], template=template) | |
| # memory = ConversationBufferMemory(input_key="question") | |
| return prompt | |
| with st.sidebar: | |
| st.subheader("Your documents") | |
| global docs | |
| docs = st.file_uploader("Upload your PDFs here and click on 'Process'", accept_multiple_files=True,type=["pdf","docx","csv","xlsx","html"]) | |
| if st.button("Process"): | |
| with st.spinner("Processing"): | |
| # raw_text = extract_text_from_pdfs(docs) | |
| # all_loaders =classify_and_load_files_into_respective_loaders(docs) | |
| # chroma_vectorstore = index_initializing_upserting_chroma_db(all_loaders) | |
| if docs: | |
| loaded_documents = load_documents_fn(docs) | |
| else: | |
| st.error("Error While loading the documents!!!Try Again!!!") | |
| if "EMBEDDINGS" not in st.session_state: | |
| EMBEDDINGS = SentenceTransformerEmbeddings(model_name=EMBEDDING_MODEL_NAME) | |
| # EMBEDDINGS = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2") | |
| st.session_state.EMBEDDINGS = EMBEDDINGS | |
| # if "DB" not in st.session_state: | |
| # DB = Chroma( | |
| # persist_directory=loaded_documents, | |
| # embedding_function=st.session_state.EMBEDDINGS, | |
| # client_settings=CHROMA_SETTINGS, | |
| # ) | |
| DB = Chroma.from_documents(loaded_documents, st.session_state.EMBEDDINGS,persist_directory="db") | |
| st.session_state.DB = DB | |
| if "RETRIEVER" not in st.session_state: | |
| RETRIEVER = DB.as_retriever() | |
| st.session_state.RETRIEVER = RETRIEVER | |
| if "LLM" not in st.session_state: | |
| LLM = load_models(model_id=MODEL_ID, model_basename=MODEL_BASENAME) | |
| # st.session_state["LLM"] = LLM | |
| if "QA" not in st.session_state: | |
| prompt = model_memory() | |
| QA = RetrievalQA.from_chain_type( | |
| llm=LLM, | |
| chain_type="stuff", | |
| retriever=RETRIEVER, | |
| return_source_documents=True, | |
| chain_type_kwargs={"prompt": prompt}, | |
| ) | |
| st.session_state["QA"] = QA | |
| st.success("LLM Initialized !!! You Chat with your documents!!") | |
| st.title('Chat With Your Documents') | |
| prompt = st.text_input('Input your prompt here') | |
| # while True: | |
| if docs is None: | |
| prompt = "" | |
| # If the user hits enter | |
| if prompt: | |
| # Then pass the prompt to the LLM | |
| response = st.session_state["QA"](prompt) | |
| answer, docs = response["result"], response["source_documents"] | |
| # ...and write it out to the screen | |
| st.write(answer) | |