Update app.py

app.py

@@ -6,13 +6,12 @@ from huggingface_hub import InferenceClient
 import re
 from datetime import datetime
 import json
-import os
 
 import arxiv
 from utils import get_md_text_abstract, search_cleaner, get_arxiv_live_search
 
-retrieve_results = 10 
-show_examples = False
+retrieve_results = 20 
+show_examples = True
 llm_models_to_choose = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'google/gemma-7b-it', 'None']
 
 generate_kwargs = dict(
@@ -22,6 +21,7 @@ generate_kwargs = dict(
     do_sample = False,
     )
 
+## RAG Model
 RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
 
 try:
@@ -32,6 +32,7 @@ try:
 except:
   gr.Warning("Retriever not working!")
 
+## Header
 mark_text = '# 🩺🔍 Search Results\n'
 header_text = "## Arxiv Paper Summary With QA Retrieval Augmented Generation \n"
 
@@ -49,6 +50,7 @@ except:
 
 database_choices = [index_info,'Arxiv Search - Latest - (EXPERIMENTAL)']
 
+## Arxiv API
 arx_client = arxiv.Client()
 is_arxiv_available = True
 check_arxiv_result = get_arxiv_live_search("What is Self Rewarding AI and how can it be used in Multi-Agent Systems?", arx_client, retrieve_results)
@@ -57,15 +59,27 @@ if len(check_arxiv_result) == 0:
   print("Arxiv search not working, switching to default search ...")
   database_choices = [index_info]
 
-if show_examples:
-    with open("sample_outputs.json", "r") as f:
-      sample_outputs = json.load(f)
-    output_placeholder = sample_outputs['output_placeholder']
-    md_text_initial = sample_outputs['search_placeholder']
-    
-else:
-    output_placeholder = None 
-    md_text_initial = ''
+
+
+## Show examples
+sample_outputs = {
+    'output_placeholder': 'The LLM will provide an answer to your question here...',
+    'search_placeholder': '''
+1. What is MoE?
+2. What are Multi Agent Systems? 
+3. What is Self Rewarding AI?
+4. What is Semantic and Episodic memory?
+5. What is AutoGen?
+6. What is ChatDev?
+7. What is Omniverse?
+8. What is Lumiere?
+9. What is SORA?
+'''
+}
+
+output_placeholder = sample_outputs['output_placeholder']
+md_text_initial = sample_outputs['search_placeholder']
+
 
 def rag_cleaner(inp):
     rank = inp['rank']
@@ -122,123 +136,87 @@ def SaveResponseAndRead(result):
     '''
     gr.HTML(documentHTML5)
 
-def save_search_results(prompt, results, response):
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    #filename = f"{timestamp}_{re.sub(r'[^\\w\\-_\\. ]', '_', prompt)}.txt"
-    filename = f"{timestamp} - {prompt}.txt"
-    with open(filename, "w") as f:
-        f.write(f"# {prompt}\n\n")
-        f.write(f"## Search Results\n\n{results}\n\n")
-        f.write(f"## LLM Response\n\n{response}\n")
-    return filename
-
-def get_past_searches():
-    txt_files = [f for f in os.listdir(".") if f.endswith(".txt") and f != "requirements.txt"]
-    return txt_files
 
 with gr.Blocks(theme = gr.themes.Soft()) as demo:
     header = gr.Markdown(header_text)
+    
+    with gr.Group():
+      msg = gr.Textbox(label = 'Search', placeholder = 'What is Mistral?')
+        
+      with gr.Accordion("Advanced Settings", open=False):
+        with gr.Row(equal_height = True):
+          llm_model = gr.Dropdown(choices = llm_models_to_choose, value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model')
+          llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results as context")
+          database_src = gr.Dropdown(choices = database_choices, value = index_info, label = 'Search Source')
+          stream_results = gr.Checkbox(value = True, label = "Stream output", visible = False)
 
-    with gr.Row():
-        with gr.Column():
-            msg = gr.Textbox(label = 'Search', placeholder = 'What is Mistral?')
-            
-            with gr.Accordion("Advanced Settings", open=False):
-                with gr.Row(equal_height = True):
-                    llm_model = gr.Dropdown(choices = llm_models_to_choose, value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model')
-                    llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results as context")
-                    database_src = gr.Dropdown(choices = database_choices, value = index_info, label = 'Search Source')
-                    stream_results = gr.Checkbox(value = True, label = "Stream output", visible = False)
-
-            output_text = gr.Textbox(show_label = True, container = True, label = 'LLM Answer', visible = True, placeholder = output_placeholder)
-            input = gr.Textbox(show_label = False, visible = False)
-            gr_md = gr.Markdown(mark_text + md_text_initial)
-
-        with gr.Column():
-            past_searches = gr.Dropdown(choices=get_past_searches(), label="Past Searches")
-            past_search_content = gr.Textbox(label="Past Search Content", visible=False)
-
-    def update_past_search_content(past_search):
-        if past_search:
-            with open(past_search, "r") as f:
-                content = f.read()
-            return gr.Textbox.update(value=content, visible=True)
-        else:
-            return gr.Textbox.update(visible=False)
+    output_text = gr.Textbox(show_label = True, container = True, label = 'LLM Answer', visible = True, placeholder = output_placeholder)
+    input = gr.Textbox(show_label = False, visible = False)
+    gr_md = gr.Markdown(mark_text + md_text_initial)
 
-    past_searches.change(update_past_search_content, past_searches, past_search_content)
-    
     def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
         prompt_text_from_data = ""
         database_to_use = database_choice
         if database_choice == index_info:
-            rag_out = get_rag(message)
+          rag_out = get_rag(message)
         else:
-            arxiv_search_success = True
-            try:
-                rag_out = get_arxiv_live_search(message, arx_client, retrieve_results)
-                if len(rag_out) == 0:
-                    arxiv_search_success = False 
-            except:
-                arxiv_search_success = False
-    
-            if not arxiv_search_success:
-                gr.Warning("Arxiv Search not working, switching to semantic search ...")
-                rag_out = get_rag(message)
-                database_to_use = index_info 
-    
+          arxiv_search_success = True
+          try:
+            rag_out = get_arxiv_live_search(message, arx_client, retrieve_results)
+            if len(rag_out) == 0:
+              arxiv_search_success = False 
+          except:
+            arxiv_search_success = False
+ 
+
+          if not arxiv_search_success:
+            gr.Warning("Arxiv Search not working, switching to semantic search ...")
+            rag_out = get_rag(message)
+            database_to_use = index_info 
+
         md_text_updated = mark_text
         for i in range(retrieve_results):
-            rag_answer = rag_out[i]
-            if i < llm_results_use:
-                md_text_paper, prompt_text = get_md_text_abstract(rag_answer, source = database_to_use, return_prompt_formatting = True)
-                prompt_text_from_data += f"{i+1}. {prompt_text}"
-            else:
-                md_text_paper = get_md_text_abstract(rag_answer, source = database_to_use)
-            md_text_updated += md_text_paper
+          rag_answer = rag_out[i]
+          if i < llm_results_use:
+            md_text_paper, prompt_text = get_md_text_abstract(rag_answer, source = database_to_use, return_prompt_formatting = True)
+            prompt_text_from_data += f"{i+1}. {prompt_text}"
+          else:
+            md_text_paper = get_md_text_abstract(rag_answer, source = database_to_use)
+          md_text_updated += md_text_paper
         prompt = get_prompt_text(message, prompt_text_from_data, llm_model_picked = llm_model_picked)
-        
-        filename = save_search_results(message, md_text_updated, "")
-        
-        with open(filename, "r") as f:
-            md_content = f.read()
-        
-        return md_content, prompt, get_past_searches()
-    
+        return md_text_updated, prompt
+
     def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2', stream_outputs = False):
-        model_disabled_text = "LLM Model is disabled"
-        output = ""
+       model_disabled_text = "LLM Model is disabled"
+       output = ""
         
-        if llm_model_picked == 'None':
-            if stream_outputs:
-                for out in model_disabled_text:
-                    output += out
-                    yield output
-                return output 
-            else:
-                return model_disabled_text
+       if llm_model_picked == 'None':
+          if stream_outputs:
+              for out in model_disabled_text:
+                output += out
+                yield output
+              return output 
+          else:
+              return model_disabled_text
               
-        client = InferenceClient(llm_model_picked)
-        try:
-            stream = client.text_generation(prompt, **generate_kwargs,  stream=stream_outputs, details=False, return_full_text=False)
+       client = InferenceClient(llm_model_picked)
+       try:
+           stream = client.text_generation(prompt, **generate_kwargs,  stream=stream_outputs, details=False, return_full_text=False)
            
-        except:
-            gr.Warning("LLM Inference rate limit reached, try again later!")
-            return ""
+       except:
+           gr.Warning("LLM Inference rate limit reached, try again later!")
+           return ""
        
-        if stream_outputs:
-            for response in stream:
-                output += response
-                SaveResponseAndRead(response)
-                yield output
-            return output
-        else:
-            return stream
+       if stream_outputs:
+           for response in stream:
+              output += response
+              SaveResponseAndRead(response)
+              yield output
+           return output
+       else:
+           return stream
 
-    msg.submit(update_with_rag_md, [msg, llm_results,  database_src, llm_model], [gr_md, input, past_searches]).success(ask_llm, [input, llm_model, stream_results], output_text).then(
-        lambda response: save_search_results(msg.value, gr_md.value, response),
-        [msg, gr_md, output_text],
-        None
-    )
+
+    msg.submit(update_with_rag_md, [msg, llm_results,  database_src, llm_model], [gr_md, input]).success(ask_llm, [input, llm_model, stream_results], output_text)
 
 demo.queue().launch()