Update app.py

app.py

@@ -6,82 +6,75 @@ 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
-import os
-import glob
 
-# 🎛️ App configuration - tweak these knobs for maximum brain power! 🧠💪
-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']
+retrieve_results = 10 
+show_examples = False
+llm_models_to_choose = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'google/gemma-2-2b-it', 'None']
+
+token = os.getenv("HF_TOKEN")
 
-# 🎭 LLM acting instructions - "To be, or not to be... verbose" 🤔
 generate_kwargs = dict(
     temperature = None,
     max_new_tokens = 512,
     top_p = None,
     do_sample = False,
-)
+    )
 
-# 🧙‍♂️ Summoning the RAG model - "Accio knowledge!" 📚✨
+## RAG Model
 RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
 
 try:
-    gr.Info("🏗️ Setting up the knowledge retriever, please wait... 🕰️")
-    rag_initial_output = RAG.search("What is Generative AI in Healthcare?", k = 1)
-    gr.Info("🎉 Retriever is up and running! Time to flex those brain muscles! 💪🧠")
+  gr.Info("Setting up retriever, please wait...")
+  rag_initial_output = RAG.search("what is Mistral?", k = 1)
+  gr.Info("Retriever working successfully!")
+    
 except:
-    gr.Warning("😱 Oh no! The retriever took a coffee break. Try again later! ☕")
+  gr.Warning("Retriever not working!")
 
-# 📜 The grand introduction - roll out the red carpet! 🎭
-mark_text = '# 🩺🔍 Search Results\n'
-header_text = "## 📚Arxiv📖Paper🔍Search - 🕵️‍♀️ Uncover, 📝 Summarize, and 🧩 Solve 🔬 Research 🤔❓ Puzzles ✍️ with 📚 Papers and 🤖 RAG AI 🧠\n"
+## Header
+mark_text = '# 🔍 Search Results\n'
+header_text = "# ArXiv CS RAG \n"
 
-# 🕰️ Time travel to find when our knowledge was last updated 🚀
 try:
-    with open("README.md", "r") as f:
-        mdfile = f.read()
-    date_pattern = r'Index Last Updated : \d{4}-\d{2}-\d{2}'
-    match = re.search(date_pattern, mdfile)
-    date = match.group().split(': ')[1]
-    formatted_date = datetime.strptime(date, '%Y-%m-%d').strftime('%d %b %Y')
-    header_text += f'Index Last Updated: {formatted_date}\n'
-    index_info = f"Semantic Search - up to {formatted_date}"  
+  with open("README.md", "r") as f:
+      mdfile = f.read()
+  date_pattern = r'Index Last Updated : \d{4}-\d{2}-\d{2}'
+  match = re.search(date_pattern, mdfile)
+  date = match.group().split(': ')[1]
+  formatted_date = datetime.strptime(date, '%Y-%m-%d').strftime('%d %b %Y')
+  header_text += f'Index Last Updated: {formatted_date}\n'
+  index_info = f"Semantic Search - up to {formatted_date}"  
 except:
-    index_info = "Semantic Search"
+  index_info = "Semantic Search"
 
-database_choices = [index_info, 'Arxiv Search - Latest - (EXPERIMENTAL)']
+database_choices = [index_info,'Arxiv Search - Latest - (EXPERIMENTAL)']
 
-# 🦉 Arxiv API - the wise old owl of academic knowledge 📜
+## 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)
+check_arxiv_result = get_arxiv_live_search("What is Mistral?", arx_client, retrieve_results)
 if len(check_arxiv_result) == 0:
-    is_arxiv_available = False
-    print("😴 Arxiv search is taking a nap, switching to default search ...")
-    database_choices = [index_info]
+  is_arxiv_available = False
+  print("Arxiv search not working, switching to default search ...")
+  database_choices = [index_info]
+
+
 
-# 🎭 Show examples - a teaser trailer for your brain! 🍿🧠
-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?
-'''
-}
+## Show examples (disabled)
+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 = ''
 
-output_placeholder = sample_outputs['output_placeholder']
-md_text_initial = sample_outputs['search_placeholder']
 
-# 🧹 Clean up the RAG output - nobody likes a messy mind! 🧼🧠
 def rag_cleaner(inp):
     rank = inp['rank']
     title = inp['document_metadata']['title']
@@ -89,20 +82,19 @@ def rag_cleaner(inp):
     date = inp['document_metadata']['_time']
     return f"{rank}. <b> {title} </b> \n Date : {date} \n Abstract: {content}"
 
-# 🎭 Craft the perfect prompt - it's showtime for the LLM! 🎬
 def get_prompt_text(question, context, formatted = True, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
     if formatted:
-        sys_instruction = f"Context:\n {context} \n Given the following scientific paper abstracts, take a deep breath and let's think step by step to answer the question. Cite the titles of your sources when answering, do not cite links or dates."
-        message = f"Question: {question}"
+      sys_instruction = f"Context:\n {context} \n Given the following scientific paper abstracts, take a deep breath and lets think step by step to answer the question. Cite the titles of your sources when answering, do not cite links or dates."
+      message = f"Question: {question}"
         
-        if 'mistralai' in llm_model_picked:
-            return f"<s>" + f"[INST] {sys_instruction}" +  f" {message}[/INST]"
-        elif 'gemma' in llm_model_picked:
-            return f"<bos><start_of_turn>user\n{sys_instruction}" +  f" {message}<end_of_turn>\n"
+      if 'mistralai' in llm_model_picked:
+          return f"<s>" + f"[INST] {sys_instruction}" +  f" {message}[/INST]"
+          
+      elif 'gemma' in llm_model_picked:
+          return f"<bos><start_of_turn>user\n{sys_instruction}" +  f" {message}<end_of_turn>\n"
           
-    return f"Context:\n {context} \n Given the following info, take a deep breath and let's think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"
+    return f"Context:\n {context} \n Given the following info, take a deep breath and lets think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"
 
-# 🕵️‍♀️ Get those juicy references - time to go treasure hunting! 💎📚
 def get_references(question, retriever, k = retrieve_results):
     rag_out = retriever.search(query=question, k=k)
     return rag_out
@@ -110,169 +102,85 @@ def get_references(question, retriever, k = retrieve_results):
 def get_rag(message):
     return get_references(message, RAG)
 
-# 🎤 Save the response and read it aloud - it's karaoke time for your brain! 🧠🎶
-def SaveResponseAndRead(result):
-    documentHTML5='''
-    <!DOCTYPE html>
-    <html>
-    <head>
-        <title>Read It Aloud</title>
-        <script type="text/javascript">
-            function readAloud() {
-                const text = document.getElementById("textArea").value;
-                const speech = new SpeechSynthesisUtterance(text);
-                window.speechSynthesis.speak(speech);
-            }
-        </script>
-    </head>
-    <body>
-        <h1>🔊 Read It Aloud</h1>
-        <textarea id="textArea" rows="10" cols="80">
-    '''
-    documentHTML5 = documentHTML5 + result
-    documentHTML5 = documentHTML5 + '''
-        </textarea>
-        <br>
-        <button onclick="readAloud()">🔊 Read Aloud</button>
-    </body>
-    </html>
-    '''
-    gr.HTML(documentHTML5)
-
-# 📁 File management functions - because even AI needs a filing system! 🗄️🤖
-
-def save_response_as_markdown(question, response):
-    timestamp = datetime.now().strftime("%Y%m%d%H%M")
-    filename = f"{timestamp}_{question[:50]}.md"  # Truncate question to 50 chars for filename
-    with open(filename, "w", encoding="utf-8") as f:
-        f.write(response)
-    return filename
-
-def list_markdown_files():
-    files = glob.glob("*.md")
-    files.sort(key=os.path.getmtime, reverse=True)
-    return [f for f in files if f != "README.md"]
-
-def delete_file(filename):
-    if filename != "README.md":
-        os.remove(filename)
-        return f"Deleted {filename}"
-    return "Cannot delete README.md"
-
-def display_markdown_contents():
-    files = list_markdown_files()
-    output = ""
-    for file in files:
-        with open(file, "r", encoding="utf-8") as f:
-            content = f.read()
-        output += f"## {file}\n\n```markdown\n{content}\n```\n\n"
-    return output
-
-# 🎨 Building the UI - it's like LEGO, but for brains! 🧠🏗️
 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 Generative AI in Healthcare?')
+      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.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.Tab("Saved Responses"):
-        refresh_button = gr.Button("🔄 Refresh File List")
-        file_list = gr.Dropdown(choices=list_markdown_files(), label="Saved Responses")
-        delete_button = gr.Button("🗑️ Delete Selected File")
-        markdown_display = gr.Markdown()
-
-    # 🔄 Update the file list - keeping things fresh! 🌿
-    def update_file_list():
-        return gr.Dropdown(choices=list_markdown_files())
-
-    refresh_button.click(update_file_list, outputs=[file_list])
-    delete_button.click(delete_file, inputs=[file_list], outputs=[markdown_display]).then(update_file_list, outputs=[file_list])
-    file_list.change(lambda x: open(x, "r", encoding="utf-8").read() if x else "", inputs=[file_list], outputs=[markdown_display])
-
-    # 🎭 The grand finale - where the magic happens! 🎩✨
     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 is taking a siesta, 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)
         return md_text_updated, prompt
 
-    # 🧠 Asking the LLM - it's like a really smart magic 8-ball! 🎱✨
     def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2', stream_outputs = False):
-        model_disabled_text = "LLM Model is taking a vacation. Try again later! 🏖️"
-        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, token = token)
+       try:
+           stream = client.text_generation(prompt, **generate_kwargs,  stream=stream_outputs, details=False, return_full_text=False)
            
-        except:
-            gr.Warning("🚦 LLM Inference hit a traffic jam! Take a breather and 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
+              yield output
+           return output
+       else:
+           return stream
 
-    # 🎬 Action! Process the query and save the response
-    def process_and_save(message, llm_results_use, database_choice, llm_model_picked):
-        md_text_updated, prompt = update_with_rag_md(message, llm_results_use, database_choice, llm_model_picked)
-        llm_response = ask_llm(prompt, llm_model_picked, stream_outputs=False)
-        full_response = f"Question: {message}\n\nResponse:\n{llm_response}\n\nReferences:\n{md_text_updated}"
-        filename = save_response_as_markdown(message, full_response)
-        return md_text_updated, prompt, llm_response, filename
 
-    # 🎬 Lights, camera, action! Let's get this show on the road! 🚀
-    msg.submit(process_and_save, [msg, llm_results, database_src, llm_model], [gr_md, input, output_text, file_list]).then(update_file_list, outputs=[file_list])
+    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)
 
-# 🎉 Launch the app - let the knowledge party begin! 🎊🧠
 demo.queue().launch()