Upload 2 files

app.py

@@ -2,165 +2,139 @@ import torch
 import transformers
 import gradio as gr
 from ragatouille import RAGPretrainedModel
-from huggingface_hub import InferenceClient
 import re
 from datetime import datetime
 import json
-
 import arxiv
-from helper import *
-
-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(
-    temperature = None,
-    max_new_tokens = 512,
-    top_p = None,
-    do_sample = False,
-    )
+from helper import rag_cleaner, get_prompt_text, get_references, get_rag, SaveResponseAndRead, get_md_text_abstract, search_cleaner, get_arxiv_live_search
+
+# Constants
+RETRIEVE_RESULTS = 20
+LLM_MODELS = ['mistralai/Mixtral-8x7B-Instruct-v0.1', 'mistralai/Mistral-7B-Instruct-v0.2', 'google/gemma-7b-it', 'None']
+DEFAULT_LLM_MODEL = 'mistralai/Mistral-7B-Instruct-v0.2'
+GENERATE_KWARGS = {
+    "temperature": None,
+    "max_new_tokens": 512,
+    "top_p": None,
+    "do_sample": False,
+}
 
-## RAG Model
+# RAG Model setup
 RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
 
 try:
-  gr.Info("Setting up retriever, please wait...")
-  rag_initial_output = RAG.search("What is Generative AI in Healthcare?", k = 1)
-  gr.Info("Retriever working successfully!")
-    
-except:
-  gr.Warning("Retriever not working!")
+    gr.Info("Setting up retriever, please wait...")
+    rag_initial_output = RAG.search("What is Generative AI in Healthcare?", k=1)
+    gr.Info("Retriever working successfully!")
+except Exception as e:
+    gr.Warning(f"Retriever not working: {str(e)}")
 
-## Header
+# Header setup
 mark_text = '# 🩺🔍 Search Results\n'
 header_text = "## Arxiv Paper Summary With QA Retrieval Augmented Generation \n"
 
 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}"  
-except:
-  index_info = "Semantic Search"
-
-database_choices = [index_info,'Arxiv Search - Latest - (EXPERIMENTAL)']
-
-## Arxiv API
+    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 FileNotFoundError:
+    index_info = "Semantic Search"
+
+database_choices = [index_info, 'Arxiv Search - Latest - (EXPERIMENTAL)']
+
+# Arxiv API setup
 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 Self Rewarding AI and how can it be used in Multi-Agent Systems?", arx_client, RETRIEVE_RESULTS)
 if len(check_arxiv_result) == 0:
-  is_arxiv_available = False
-  print("Arxiv search not working, 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
-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']
-
-
-with gr.Blocks(theme = gr.themes.Soft()) as demo:
+# Gradio UI setup
+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?')
+        search_query = gr.Textbox(label='Search', placeholder='What is Generative AI in Healthcare?')
         
-      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)
-
-    def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
+        with gr.Accordion("Advanced Settings", open=False):
+            with gr.Row(equal_height=True):
+                llm_model = gr.Dropdown(choices=LLM_MODELS, value=DEFAULT_LLM_MODEL, 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)
+    input = gr.Textbox(show_label=False, visible=False)
+    gr_md = gr.Markdown(mark_text)
+    
+    def update_with_rag_md(search_query, llm_results_use=5, database_choice=index_info, llm_model_picked=DEFAULT_LLM_MODEL):
         prompt_text_from_data = ""
         database_to_use = database_choice
+        
         if database_choice == index_info:
-          rag_out = get_rag(message)
+            rag_out = get_rag(search_query, RAG, RETRIEVE_RESULTS)
         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(search_query, arx_client, RETRIEVE_RESULTS)
+                if len(rag_out) == 0:
+                    arxiv_search_success = False
+            except Exception as e:
+                arxiv_search_success = False
+                gr.Warning(f"Arxiv Search not working: {str(e)}, switching to semantic search ...")
+            
+            if not arxiv_search_success:
+                rag_out = get_rag(search_query, RAG, RETRIEVE_RESULTS)
+                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
-        prompt = get_prompt_text(message, prompt_text_from_data, llm_model_picked = llm_model_picked)
+        for i, rag_answer in enumerate(rag_out):
+            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(search_query, prompt_text_from_data, llm_model_picked=llm_model_picked)
         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 = ""
+    
+    def ask_llm(prompt, llm_model_picked=DEFAULT_LLM_MODEL, stream_outputs=False):
+        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
-              
-       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 ""
-       
-       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]).success(ask_llm, [input, llm_model, stream_results], output_text)
+        if llm_model_picked == 'None':
+            if stream_outputs:
+                for out in model_disabled_text:
+                    output += out
+                    yield output
+            else:
+                return model_disabled_text
+        
+        client = InferenceClient(llm_model_picked)
+        try:
+            response = client.text_generation(prompt, stream=stream_outputs, details=False, return_full_text=False, **GENERATE_KWARGS)
+            
+            if stream_outputs:
+                for token in response:
+                    output += token
+                    yield SaveResponseAndRead(output)
+            else:
+                output = response
+        except Exception as e:
+            gr.Warning(f"LLM Inference failed: {str(e)}")
+            output = ""
+        
+        return output
+    
+    search_query.submit(update_with_rag_md, [search_query, llm_results, database_src, llm_model], [gr_md, input]).success(ask_llm, [input, llm_model, stream_results], output_text)
 
 demo.queue().launch()

helper.py

@@ -1,3 +1,6 @@
+import sys
+import gradio as gr
+from huggingface_hub import InferenceClient
 import datetime 
 import string
 import nltk
@@ -5,11 +8,6 @@ nltk.download('stopwords')
 from nltk.corpus import stopwords
 stop_words = stopwords.words('english')
 import arxiv
-import gradio as gr
-import re
-from datetime import datetime
-import json
-
 
 def rag_cleaner(inp):
     rank = inp['rank']
@@ -18,28 +16,27 @@ def rag_cleaner(inp):
     date = inp['document_metadata']['_time']
     return f"{rank}. <b> {title} </b> \n Date : {date} \n Abstract: {content}"
 
-def get_prompt_text(question, context, formatted = True, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
+def get_prompt_text(question, context, formatted=True, llm_model_picked='mistralai/Mistral-7B-Instruct-v0.2'):
+    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 formatted:
-      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>[INST] {sys_instruction} {message}[/INST]"
+        elif 'gemma' in llm_model_picked:
+            return f"<bos><start_of_turn>user\n{sys_instruction} {message}<end_of_turn>\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"
 
-def get_references(question, retriever, k = retrieve_results):
+def get_references(question, retriever, k):
     rag_out = retriever.search(query=question, k=k)
     return rag_out
 
-def get_rag(message):
-    return get_references(message, RAG)
+def get_rag(message, RAG, RETRIEVE_RESULTS):
+    return get_references(message, RAG, k=RETRIEVE_RESULTS)
 
 def SaveResponseAndRead(result):
-    documentHTML5='''
+    documentHTML5 = '''
     <!DOCTYPE html>
     <html>
     <head>
@@ -56,17 +53,15 @@ def SaveResponseAndRead(result):
         <h1>🔊 Read It Aloud</h1>
         <textarea id="textArea" rows="10" cols="80">
     '''
-    documentHTML5 = documentHTML5 + result
-    documentHTML5 = documentHTML5 + '''
+    documentHTML5 += result
+    documentHTML5 += '''
         </textarea>
         <br>
         <button onclick="readAloud()">🔊 Read Aloud</button>
     </body>
     </html>
     '''
-    gr.HTML(documentHTML5)
-
-
+    return gr.HTML(documentHTML5)
 
 
 def get_md_text_abstract(rag_answer, source = ['Arxiv Search', 'Semantic Search'][1], return_prompt_formatting = False):