Somnath3570/PDF_based_knowledge_management_system

1import os2import streamlit as st3# Update these imports4from langchain_community.embeddings import HuggingFaceEmbeddings5from langchain.chains import RetrievalQA6from langchain_community.vectorstores import FAISS7from langchain_core.prompts import PromptTemplate8from langchain_huggingface import HuggingFaceEndpoint9 10from dotenv import load_dotenv, find_dotenv11load_dotenv(find_dotenv())12 13DB_FAISS_PATH = "vectorstore/db_faiss"14 15@st.cache_resource16def get_vectorstore():17    embedding_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')18    db = FAISS.load_local(DB_FAISS_PATH, embedding_model, allow_dangerous_deserialization=True)19    return db20 21def set_custom_prompt(custom_prompt_template):22    prompt = PromptTemplate(template=custom_prompt_template, input_variables=["context", "question"])23    return prompt24 25def load_llm(huggingface_repo_id, HF_TOKEN):26    llm = HuggingFaceEndpoint(27        repo_id=huggingface_repo_id,28        task="text-generation",  # Add this line29        temperature=0.5,30        model_kwargs={31            "token": HF_TOKEN,32            "max_length": 512  # Changed to integer33        }34    )35    return llm36 37def main():38    st.title("Forecasting discharge outcomes for critically ILL patients using machine learning!")39    40    if 'messages' not in st.session_state:41        st.session_state.messages = []42    43    for message in st.session_state.messages:44        st.chat_message(message['role']).markdown(message['content'])45    46    prompt = st.chat_input("Pass your prompt here")47    48    if prompt:49        st.chat_message('user').markdown(prompt)50        st.session_state.messages.append({'role': 'user', 'content': prompt})51        52        CUSTOM_PROMPT_TEMPLATE = """53        Use the pieces of information provided in the context to answer user's question.54        If you dont know the answer, just say that you dont know, dont try to make up an answer.55        56        Dont provide anything out of the given context57        58        Context: {context}59        Question: {question}60        61        Start the answer directly. No small talk please.62        """63        64        HUGGINGFACE_REPO_ID = "mistralai/Mistral-7B-Instruct-v0.3"65        HF_TOKEN = os.environ.get("HF_TOKEN")66        67        try:68            with st.spinner("Thinking..."):  # Add loading indicator69                vectorstore = get_vectorstore()70                if vectorstore is None:71                    st.error("Failed to load the vector store")72                    return73 74                qa_chain = RetrievalQA.from_chain_type(75                    llm=load_llm(huggingface_repo_id=HUGGINGFACE_REPO_ID, HF_TOKEN=HF_TOKEN),76                    chain_type="stuff",77                    retriever=vectorstore.as_retriever(search_kwargs={'k': 3}),78                    return_source_documents=True,79                    chain_type_kwargs={'prompt': set_custom_prompt(CUSTOM_PROMPT_TEMPLATE)}80                )81                82                response = qa_chain.invoke({'query': prompt})83                84                result = response["result"]85                source_documents = response["source_documents"]86                87                # Format source documents more cleanly88                source_docs_text = "\n\n**Source Documents:**\n"89                for i, doc in enumerate(source_documents, 1):90                    source_docs_text += f"{i}. Page {doc.metadata.get('page', 'N/A')}: {doc.page_content[:200]}...\n\n"91                92                result_to_show = f"{result}\n{source_docs_text}"93                94                st.chat_message('assistant').markdown(result_to_show)95                st.session_state.messages.append({'role': 'assistant', 'content': result_to_show})96                97        except Exception as e:98            st.error(f"Error: {str(e)}")99            st.error("Please check your HuggingFace token and model access permissions")100 101if __name__ == "__main__":102    main()