Initialise the app

app.py

@@ -0,0 +1,134 @@
+"""
+Streamlit app
+"""
+try:
+    __import__('pysqlite3')
+    import sys
+    sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
+except:
+    pass
+
+import os
+import streamlit as st
+from ragvizexpander import RAGVizChain
+from ragvizexpander.llms import *
+from ragvizexpander.embeddings import *
+from ragvizexpander.splitters import RecursiveChar2TokenSplitter
+
+st.set_page_config(
+    page_title="RAGVizExpander Demo",
+    page_icon="🔬",
+    layout="wide"
+)
+
+os.environ['OPENAI_API_KEY'] = st.secrets["OPENAI_API_KEY"]
+os.environ['HF_API_KEY'] = st.secrets["HF_API_KEY"]
+
+
+if "chart" not in st.session_state:
+    st.session_state['chart'] = None
+
+if "loaded" not in st.session_state:
+    st.session_state['loaded'] = False
+
+st.title("RAGVizExpander Demo🔬")
+st.markdown("📦 More details can be found at the GitHub repo [here](https://github.com/KKenny0/RAGVizExpander)")
+
+if not st.session_state['loaded']:
+    main_page = st.empty()
+    main_button = st.empty()
+    with main_page.container():
+        uploaded_file = st.file_uploader("Upload your file",
+                                         label_visibility="collapsed",
+                                         type=['pdf', 'docx', 'txt', 'pptx'])
+
+        # --- setting llm model
+        st.markdown("### Settings for *LLM* model")
+
+        st.session_state["llm_model_type"] = st.radio("Select type of llm model",
+                                                      ["OpenAI", "Ollama"],
+                                                      horizontal=True)
+
+        if st.session_state["llm_model_type"] == "OpenAI":
+            st.session_state["openai_llm_base_url"] = st.text_input("Enter OpenAI LLM API Base")
+            st.session_state["openai_llm_api_key"] = st.text_input("Enter OpenAI LLM API Key")
+            st.session_state["openai_llm_model"] = st.text_input("Enter OpenAI LLM model name")
+            st.session_state["chosen_llm_model"] = ChatOpenAI(
+                base_url=st.session_state["openai_llm_base_url"],
+                api_key=st.session_state["openai_llm_api_key"],
+                model_name=st.session_state["openai_llm_model"],
+            )
+        else:
+            st.session_state["ollama_llm_model"] = st.text_input("Enter Ollama model name")
+            st.session_state["chosen_llm_model"] = ChatOllama(model_name=st.session_state["ollama_llm_model"])
+
+        st.markdown("""---""")
+
+        # --- setting embedding model
+        st.markdown("### Settings for *EMBEDDING* model")
+
+        st.session_state["embedding_model_type"] = st.radio("Select type of embedding model",
+                                                            ["OpenAI", "SentenceTransformer", "HuggingFace", "TEI"],
+                                                            horizontal=True)
+
+        if st.session_state["embedding_model_type"] == "OpenAI":
+            st.session_state["openai_embed_model"] = st.selectbox("Select embedding model",
+                                                                      ["text-embedding-3-small",
+                                                                       "text-embedding-3-large",
+                                                                       "text-embedding-ada-002"])
+            st.session_state["openai_embed_api_key"] = st.text_input("Enter OpenAI Embedding API Key")
+            st.session_state["openai_embed_api_base"] = st.text_input("Enter OpenAI Embedding API Base")
+            st.session_state["chosen_embedding_model"] = OpenAIEmbeddings(
+                api_base=st.session_state["openai_embed_api_base"],
+                api_key=st.session_state["openai_embed_api_key"],
+                model_name=st.session_state["openai_embed_model"],
+            )
+
+        elif st.session_state["embedding_model_type"] == "HuggingFace":
+            st.session_state["hf_embed_model"] = st.text_input("Enter HF repository name")
+            st.session_state["hf_api_key"] = st.text_input("Enter HF API key")
+            st.session_state["chosen_embedding_model"] = HuggingFaceEmbeddings(
+                model_name=st.session_state["hf_embed_model"],
+                api_key=st.session_state["hf_api_key"]
+            )
+
+        else:
+            st.session_state["tei_api_url"] = st.text_input("Enter TEI(Text-Embedding-Inference) api url")
+            st.session_state["chosen_embedding_model"] = TEIEmbeddings(
+                api_url=st.session_state["tei_api_url"]
+            )
+
+        st.markdown("""---""")
+
+        # --- setting chunking parameters
+        st.markdown("### Settings for *CHUNKING* model")
+        st.session_state["chunk_size"] = st.number_input("Chunk size", value=500, min_value=100, max_value=1000, step=100)
+        st.session_state["chunk_overlap"] = st.number_input("Chunk overlap", value=0, min_value=0, max_value=100, step=10)
+        st.session_state["split_func"] = RecursiveChar2TokenSplitter(
+            chunk_size=st.session_state["chunk_size"],
+            chunk_overlap=st.session_state["chunk_overlap"],
+        )
+
+    if st.button("Build Vector DB"):
+        st.session_state["client"] = RAGVizChain(embedding_model=st.session_state["chosen_embedding_model"],
+                                                 llm=st.session_state["chosen_llm_model"],
+                                                 split_func=st.session_state["split_func"])
+        main_page.empty()
+        main_button.empty()
+        with st.spinner("Building Vector DB"):
+            st.session_state["client"].load_data(uploaded_file,)
+            st.session_state['loaded'] = True
+            st.rerun()
+else:
+    col1, col2 = st.columns(2)
+    st.session_state['query'] = col1.text_area("Enter your query here")
+    st.session_state['technique'] = col1.radio("Select retrival technique", ["naive", "HyAE", "multi_qns"], horizontal=True)
+    st.session_state['top_k'] = col1.number_input("Top k", value=5, min_value=1, max_value=10, step=1)
+    if col1.button("Execute Query"):
+            st.session_state['chart'] = st.session_state["client"].visualize_query(st.session_state['query'], retrieval_method=st.session_state['technique'], top_k=st.session_state['top_k'])
+    if st.session_state['chart'] is not None:
+        col2.plotly_chart(st.session_state['chart'])
+
+    if col1.button("Reset Application"):
+        st.session_state['loaded'] = False
+        st.rerun()