xnetba/ChatPDF
0
1import base642import streamlit as st3from streamlit_chat import message4from streamlit_extras.colored_header import colored_header5from backend import QnASystem6from schema import TransformType, EmbeddingTypes, IndexerType, BotType7 8kwargs = {}9source_docs = []10st.set_page_config(page_title="PDFChat - An LLM-powered experimentation app")11 12if "qna_system" not in st.session_state:13 st.session_state.qna_system = QnASystem()14 15def show_pdf(f):16 f.seek(0)17 base64_pdf = base64.b64encode(f.read()).decode('utf-8')18 pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="700" height="800" ' \19 f'type="application/pdf"></iframe>'20 st.markdown(pdf_display, unsafe_allow_html=True)21 22def model_settings():23 kwargs["temperature"] = st.slider("Temperature", max_value=1.0, min_value=0.0)24 kwargs["max_tokens"] = st.number_input("Max Token", min_value=0, value=512)25 26st.title("PDF Question and Answering")27 28tab1, tab2, tab3 = st.tabs(["Upload and Ingest PDF", "Ask", "Show PDF"])29 30with st.sidebar:31 st.header("Advance Setting ⚙️")32 require_pdf = st.checkbox("Show PDF", value=1)33 st.markdown('---')34 kwargs["bot_type"] = st.selectbox("Bot Type", options=BotType)35 st.markdown("---")36 st.text("Model Parameters")37 kwargs["return_documents"] = st.checkbox("Require Source Documents", value=True)38 text_transform = st.selectbox("Text Transformer", options=TransformType)39 st.markdown("---")40 selected_model = st.selectbox("Select Model", options=EmbeddingTypes)41 match selected_model:42 case EmbeddingTypes.OPENAI:43 api_key = st.text_input("OpenAI API Key", placeholder="sk-...", type="password")44 if not api_key.startswith('sk-'):45 st.warning('Please enter your OpenAI API key!', icon='⚠')46 model_settings()47 case EmbeddingTypes.HUGGING_FACE:48 api_key = st.text_input("Hugging Face API Key", placeholder="hg-...", type="password")49 if not api_key.startswith('hg-'):50 st.warning('Please enter your HuggingFace API key!', icon='⚠')51 kwargs["model_name"] = st.selectbox("Choose Model", options=["google/flan-t5-xxl"])52 model_settings()53 case EmbeddingTypes.COHERE:54 api_key = st.text_input("Cohere API Key", placeholder="...", type="password")55 if not api_key:56 st.warning('Please enter your Cohere API key!', icon='⚠')57 model_settings()58 case _:59 api_key = None60 kwargs["api_key"] = api_key61 st.markdown("---")62 63 vector_indexer = st.selectbox("Vector Indexer", options=IndexerType)64 match vector_indexer:65 case IndexerType.ELASTICSEARCH:66 kwargs["elasticsearch_url"] = st.text_input("Elastic Search URL: ")67 if not kwargs.get("elasticsearch_url"):68 st.warning("Please enter your elastic search url", icon='⚠')69 kwargs["elasticsearch_index"] = st.text_input("Elastic Search Index: ")70 if not kwargs.get("elasticsearch_index"):71 st.warning("Please enter your elastic search index", icon='⚠')72 73 st.markdown("---")74 st.text("Chain Settings")75 kwargs["chain_type"] = st.selectbox("Chain Type", options=["stuff", "map_reduce"])76 kwargs["search_type"] = st.selectbox("Search Type", options=["similarity"])77 st.markdown("---")78 79with tab1:80 uploaded_file = st.file_uploader("Upload and Ingest PDF 🚀", type="pdf")81 if uploaded_file:82 with st.spinner("Uploading and Ingesting"):83 documents = st.session_state.qna_system.read_and_load_pdf(uploaded_file)84 if selected_model == EmbeddingTypes.NA:85 st.warning("Please select the model", icon='⚠')86 else:87 st.session_state.qna_system.build_chain(transform_type=text_transform, embedding_type=selected_model,88 indexer_type=vector_indexer, **kwargs)89 90def generate_response(prompt):91 if prompt and uploaded_file:92 response = st.session_state.qna_system.ask_question(prompt)93 return response.get("answer", response.get("result", "")), response.get("source_documents")94 return "", []95 96with tab2:97 if not uploaded_file:98 st.warning("Please upload PDF", icon='⚠')99 else:100 match kwargs["bot_type"]:101 case BotType.qna:102 with st.container():103 with st.form('my_form'):104 text = st.text_area("", placeholder='Ask me...')105 submitted = st.form_submit_button('Submit')106 if text:107 st.write(f"Question:\n{text}")108 response, source_docs = generate_response(text)109 st.write(response)110 case BotType.conversational:111 # Generate empty lists for generated and past.112 ## generated stores AI generated responses113 if 'generated' not in st.session_state:114 st.session_state['generated'] = ["Hi! I'm PDF Assistant 🤖, How may I help you?"]115 ## past stores User's questions116 if 'past' not in st.session_state:117 st.session_state['past'] = ['Hi!']118 119 input_container = st.container()120 colored_header(label='', description='', color_name='blue-30')121 response_container = st.container()122 response = ""123 124 def get_text():125 input_text = st.text_input("You: ", "", key="input")126 return input_text127 128 with input_container:129 user_input = get_text()130 if st.button("Clear"):131 st.session_state.generated.clear()132 st.session_state.past.clear()133 134 with response_container:135 if user_input:136 response, source_docs = generate_response(user_input)137 st.session_state.past.append(user_input)138 st.session_state.generated.append(response)139 140 if st.session_state['generated']:141 for i in range(len(st.session_state['generated'])):142 message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')143 message(st.session_state["generated"][i], key=str(i))144 145 require_document = st.container()146 if kwargs["return_documents"]:147 with require_document:148 with st.expander("Related Documents", expanded=False):149 for source in source_docs:150 metadata = source.metadata151 st.write("{source} - {page_no}".format(source=metadata.get("source"),152 page_no=metadata.get("page_no")))153 st.write(source.page_content)154 st.markdown("---")155 156with tab3:157 if require_pdf and uploaded_file:158 show_pdf(uploaded_file)159 elif uploaded_file:160 st.warning("Feature not enabled.", icon='⚠')161 else:162 st.warning("Please upload PDF", icon='⚠')163