Spaces:
Sleeping
Sleeping
| """ | |
| 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() | |