Update app.py

app.py

@@ -59,191 +59,199 @@ class GeminiTokens:
 def main():
     with st.sidebar:
         st.title('Document Summarization and QA System')
-
-        # Select Provider
-        provider = st.selectbox(
-            label="Select LLM Provider",
-            options=['google', 'huggingface', 'mistralai', 'openai'],
-            index=3
-        )
-
-        # Select LLM
-        if provider == 'google':
-            llm_list = ['gemini-1.0-pro', 'gemini-1.5-flash', 'gemini-1.5-pro']
-        elif provider == 'huggingface':
-            llm_list = []
-        elif provider == 'mistralai':
-            llm_list = ["mistral-large-latest", "open-mistral-nemo-latest"]
-        elif provider == 'openai':
-            llm_list = ['gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'gpt-4o', 'gpt-4o-mini']
-        else:
-            llm_list = []
-
-        if provider == 'huggingface':
-            llm_name = st.text_input(
-                "Enter LLM namespace/model-name",
-                value="HuggingFaceH4/zephyr-7b-alpha",
-            )
-
-            # Also give the user the option for different embedding models, too
-            embed_name = st.text_input(
-                label="Enter embedding namespace/model-name",
-                value="BAAI/bge-small-en-v1.5",
-            )
-        else:
-            llm_name = st.selectbox(
-                label="Select LLM Model",
-                options=llm_list,
-                index=0
+        
+        with st.form(key="model_settings"):
+            # Select Provider
+            provider = st.selectbox(
+                label="Select LLM Provider",
+                options=['google', 'huggingface', 'mistralai', 'openai'],
+                index=3
             )
-
-        # Temperature
-        temperature = st.slider(
-            "Temperature",
-            min_value=0.0, 
-            max_value=1.0, 
-            value=0.0, 
-            step=0.05, 
-        )
-
-        # Enter Parsing API Key
-        parse_key = st.text_input(
-            "Enter your LlamaParse API Key",
-            value=None
-        )
-
-        # Enter LLM API Key
-        llm_key = st.text_input(
-            "Enter your LLM provider API Key",
-            value=None,
-        )
-
-        # Create LLM
-        # Global tokenization needs to be consistent with LLM for token counting
-        # https://docs.llamaindex.ai/en/stable/module_guides/models/llms/
-        if llm_key is not None:
+    
+            # Select LLM
             if provider == 'google':
-                from llama_index.llms.gemini import Gemini
-                from llama_index.embeddings.gemini import GeminiEmbedding
-                max_output_tokens = 8192 # https://firebase.google.com/docs/vertex-ai/gemini-models
-                
-                os.environ['GOOGLE_API_KEY'] = str(llm_key)
-                Settings.llm = Gemini(
-                    model=f"models/{llm_name}",
-                    token=os.environ.get("GOOGLE_API_KEY"),
-                    temperature=temperature,
-                    max_tokens=max_output_tokens
+                llm_list = ['gemini-1.0-pro', 'gemini-1.5-flash', 'gemini-1.5-pro']
+            elif provider == 'huggingface':
+                llm_list = []
+            elif provider == 'mistralai':
+                llm_list = ["mistral-large-latest", "open-mistral-nemo-latest"]
+            elif provider == 'openai':
+                llm_list = ['gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'gpt-4o', 'gpt-4o-mini']
+            else:
+                llm_list = []
+    
+            if provider == 'huggingface':
+                llm_name = st.text_input(
+                    "Enter LLM namespace/model-name",
+                    value="HuggingFaceH4/zephyr-7b-alpha",
+                )
+    
+                # Also give the user the option for different embedding models, too
+                embed_name = st.text_input(
+                    label="Enter embedding namespace/model-name",
+                    value="BAAI/bge-small-en-v1.5",
                 )
-                Settings.tokenizer = GeminiTokens(llm_name) 
-                Settings.num_output = max_output_tokens
-                Settings.embed_model = GeminiEmbedding(
-                    model_name="models/text-embedding-004", api_key=os.environ.get("GOOGLE_API_KEY") #, title="this is a document"
+            else:
+                llm_name = st.selectbox(
+                    label="Select LLM Model",
+                    options=llm_list,
+                    index=0
                 )
-                if llm_name == 'gemini-1.0-pro':
-                    total_token_limit = 32760
-                else:
-                    total_token_limit = 1e6
-                Settings.context_window = total_token_limit -  max_output_tokens # Gemini counts total tokens
-            elif provider == 'huggingface':
-                if llm_name is not None and embed_name is not None:
-                    from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI 
-                    from llama_index.embeddings.huggingface import HuggingFaceInferenceAPIEmbedding
-                    from transformers import AutoTokenizer
-
-                    max_output_tokens = 2048 # Just a generic value
+    
+            # Temperature
+            temperature = st.slider(
+                "Temperature",
+                min_value=0.0, 
+                max_value=1.0, 
+                value=0.0, 
+                step=0.05, 
+            )
 
-                    os.environ['HF_TOKEN'] = str(llm_key)
-                    Settings.llm = HuggingFaceInferenceAPI(
-                        model_name=llm_name, 
-                        token=os.environ.get("HF_TOKEN"),
+            similarity_top_k = st.number_input("Top k nodes to retrieve (similarity_top_k)", min_value=1, max_value=100, value=5, step=1)
+            similarity_cutoff = st.slider("Select node similarity cutoff", min_value=0.0, max_value=1.0, value=0.7)
+            
+            # Enter Parsing API Key
+            parse_key = st.text_input(
+                "Enter your LlamaParse API Key",
+                value=None
+            )
+    
+            # Enter LLM API Key
+            llm_key = st.text_input(
+                "Enter your LLM provider API Key",
+                value=None,
+            )
+    
+            # Create LLM
+            # Global tokenization needs to be consistent with LLM for token counting
+            # https://docs.llamaindex.ai/en/stable/module_guides/models/llms/
+            if llm_key is not None:
+                if provider == 'google':
+                    from llama_index.llms.gemini import Gemini
+                    from llama_index.embeddings.gemini import GeminiEmbedding
+                    max_output_tokens = 8192 # https://firebase.google.com/docs/vertex-ai/gemini-models
+                    
+                    os.environ['GOOGLE_API_KEY'] = str(llm_key)
+                    Settings.llm = Gemini(
+                        model=f"models/{llm_name}",
+                        token=os.environ.get("GOOGLE_API_KEY"),
                         temperature=temperature,
                         max_tokens=max_output_tokens
                     )
-                    Settings.tokenizer = AutoTokenizer.from_pretrained(
-                        llm_name,
-                        token=os.environ.get("HF_TOKEN"),
+                    Settings.tokenizer = GeminiTokens(llm_name) 
+                    Settings.num_output = max_output_tokens
+                    Settings.embed_model = GeminiEmbedding(
+                        model_name="models/text-embedding-004", api_key=os.environ.get("GOOGLE_API_KEY") #, title="this is a document"
+                    )
+                    if llm_name == 'gemini-1.0-pro':
+                        total_token_limit = 32760
+                    else:
+                        total_token_limit = 1e6
+                    Settings.context_window = total_token_limit -  max_output_tokens # Gemini counts total tokens
+                elif provider == 'huggingface':
+                    if llm_name is not None and embed_name is not None:
+                        from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI 
+                        from llama_index.embeddings.huggingface import HuggingFaceInferenceAPIEmbedding
+                        from transformers import AutoTokenizer
+    
+                        max_output_tokens = 2048 # Just a generic value
+    
+                        os.environ['HF_TOKEN'] = str(llm_key)
+                        Settings.llm = HuggingFaceInferenceAPI(
+                            model_name=llm_name, 
+                            token=os.environ.get("HF_TOKEN"),
+                            temperature=temperature,
+                            max_tokens=max_output_tokens
+                        )
+                        Settings.tokenizer = AutoTokenizer.from_pretrained(
+                            llm_name,
+                            token=os.environ.get("HF_TOKEN"),
+                        )
+                        Settings.num_output = max_output_tokens
+                        Settings.embed_model = HuggingFaceInferenceAPIEmbedding(
+                            model_name=embed_name
+                        )
+                        Settings.context_window = 4096 # Just a generic value
+                elif provider == 'mistralai':
+                    from llama_index.llms.mistralai import MistralAI
+                    from llama_index.embeddings.mistralai import MistralAIEmbedding
+                    max_output_tokens = 8192 # Based on internet consensus since this is not well documented
+                    
+                    os.environ['MISTRAL_API_KEY'] = str(llm_key)
+                    Settings.llm = MistralAI(
+                        model=llm_name, 
+                        temperature=temperature,
+                        max_tokens=max_output_tokens,
+                        random_seed=42,
+                        safe_mode=True
                     )
+                    Settings.tokenizer = MistralTokens(llm_name)
                     Settings.num_output = max_output_tokens
-                    Settings.embed_model = HuggingFaceInferenceAPIEmbedding(
-                        model_name=embed_name
+                    Settings.embed_model = MistralAIEmbedding(
+                        model_name="mistral-embed", 
+                        api_key=os.environ.get("MISTRAL_API_KEY")
                     )
-                    Settings.context_window = 4096 # Just a generic value
-            elif provider == 'mistralai':
-                from llama_index.llms.mistralai import MistralAI
-                from llama_index.embeddings.mistralai import MistralAIEmbedding
-                max_output_tokens = 8192 # Based on internet consensus since this is not well documented
-                
-                os.environ['MISTRAL_API_KEY'] = str(llm_key)
-                Settings.llm = MistralAI(
-                    model=llm_name, 
-                    temperature=temperature,
-                    max_tokens=max_output_tokens,
-                    random_seed=42,
-                    safe_mode=True
-                )
-                Settings.tokenizer = MistralTokens(llm_name)
-                Settings.num_output = max_output_tokens
-                Settings.embed_model = MistralAIEmbedding(
-                    model_name="mistral-embed", 
-                    api_key=os.environ.get("MISTRAL_API_KEY")
-                )
-                Settings.context_window = 128000 # 128k for flagship models - doesn't seem to count input tokens
-            elif provider == 'openai':
-                from llama_index.llms.openai import OpenAI
-                from llama_index.embeddings.openai import OpenAIEmbedding
-
-                # https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4
-                if llm_name == 'gpt-3.5-turbo':
-                    max_output_tokens = 4096
-                    context_window = 16385
-                elif llm_name == 'gpt-4':
-                    max_output_tokens = 8192
-                    context_window = 8192
-                elif llm_name == 'gpt-4-turbo':
-                    max_output_tokens = 4096
-                    context_window = 128000
-                elif llm_name == 'gpt-4o':
-                    max_output_tokens = 4096
-                    context_window = 128000
-                elif llm_name == 'gpt-4o-mini':
-                    max_output_tokens = 16384
-                    context_window = 128000
-
-                os.environ["OPENAI_API_KEY"] = str(llm_key)
-                Settings.llm = OpenAI(
-                    model=llm_name, 
-                    temperature=temperature,
-                    max_tokens=max_output_tokens
+                    Settings.context_window = 128000 # 128k for flagship models - doesn't seem to count input tokens
+                elif provider == 'openai':
+                    from llama_index.llms.openai import OpenAI
+                    from llama_index.embeddings.openai import OpenAIEmbedding
+    
+                    # https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4
+                    if llm_name == 'gpt-3.5-turbo':
+                        max_output_tokens = 4096
+                        context_window = 16385
+                    elif llm_name == 'gpt-4':
+                        max_output_tokens = 8192
+                        context_window = 8192
+                    elif llm_name == 'gpt-4-turbo':
+                        max_output_tokens = 4096
+                        context_window = 128000
+                    elif llm_name == 'gpt-4o':
+                        max_output_tokens = 4096
+                        context_window = 128000
+                    elif llm_name == 'gpt-4o-mini':
+                        max_output_tokens = 16384
+                        context_window = 128000
+    
+                    os.environ["OPENAI_API_KEY"] = str(llm_key)
+                    Settings.llm = OpenAI(
+                        model=llm_name, 
+                        temperature=temperature,
+                        max_tokens=max_output_tokens
+                    )
+                    Settings.tokenizer = tiktoken.encoding_for_model(llm_name).encode
+                    Settings.num_output = max_output_tokens
+                    Settings.embed_model = OpenAIEmbedding()
+                    Settings.context_window = context_window 
+                else:
+                    raise NotImplementedError(f"{provider} is not supported yet")
+    
+            uploaded_file = st.file_uploader(
+                "Choose a PDF file to upload", 
+                type=['pdf'], 
+                accept_multiple_files=False
+            )
+    
+            parsed_document = None
+            if uploaded_file is not None:
+                # Parse the file
+                parser = LlamaParse(
+                    api_key=parse_key,  # Can also be set in your env as LLAMA_CLOUD_API_KEY
+                    result_type="text"  # "markdown" and "text" are available
                 )
-                Settings.tokenizer = tiktoken.encoding_for_model(llm_name).encode
-                Settings.num_output = max_output_tokens
-                Settings.embed_model = OpenAIEmbedding()
-                Settings.context_window = context_window 
-            else:
-                raise NotImplementedError(f"{provider} is not supported yet")
-
-        uploaded_file = st.file_uploader(
-            "Choose a PDF file to upload", 
-            type=['pdf'], 
-            accept_multiple_files=False
-        )
-
-        parsed_document = None
-        if uploaded_file is not None:
-            # Parse the file
-            parser = LlamaParse(
-                api_key=parse_key,  # Can also be set in your env as LLAMA_CLOUD_API_KEY
-                result_type="text"  # "markdown" and "text" are available
+    
+                # Create a temporary directory to save the file then load and parse it
+                temp_dir = tempfile.TemporaryDirectory()
+                temp_filename = os.path.join(temp_dir.name, uploaded_file.name)
+                with open(temp_filename, "wb") as f:
+                    f.write(uploaded_file.getvalue())
+                parsed_document = parser.load_data(temp_filename)
+                temp_dir.cleanup()
+            
+            submit_button = st.form_submit_button(
+                "Construct RAG"
             )
 
-            # Create a temporary directory to save the file then load and parse it
-            temp_dir = tempfile.TemporaryDirectory()
-            temp_filename = os.path.join(temp_dir.name, uploaded_file.name)
-            with open(temp_filename, "wb") as f:
-                f.write(uploaded_file.getvalue())
-            parsed_document = parser.load_data(temp_filename)
-            temp_dir.cleanup()
-
     col1, col2 = st.columns(2)
 
     with col2:
@@ -281,9 +289,15 @@ def main():
 
         run = st.button("Answer", type="primary")
 
-        if parsed_document is not None and run:
+        if parsed_document is not None and run and submit_button:
             index = VectorStoreIndex.from_documents(parsed_document)
-            query_engine = index.as_query_engine()
+            query_engine = index.as_query_engine(
+                similarity_top_k=similarity_top_k,   
+                similarity_cutoff=similarity_cutoff,
+                response_mode='compact',
+                # text_qa_template=text_qa_template,
+                # refine_template=refine_template,
+            )
             response = query_engine.query(prompt)
             st.write(response.response)