Added option for selecting LLM and number of abstracts as input

app.py

@@ -5,9 +5,6 @@ from ragatouille import RAGPretrainedModel
 from huggingface_hub import InferenceClient
 
 retrieve_results = 10 
-llm_results = 5 
-
-client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 generate_kwargs = dict(
     temperature = None,
@@ -24,7 +21,7 @@ try:
 except:
   gr.Warning("Retriever not working!")
 
-mark_text = '# 📚 Search Results\n'
+mark_text = '# 🔍 Search Results\n'
 
 def rag_cleaner(inp):
     rank = inp['rank']
@@ -47,29 +44,43 @@ def get_rag(message):
         return get_references(message, RAG)
 
 with gr.Blocks(theme = gr.themes.Soft()) as demo:
+    header = gr.Markdown("# ArXiv RAG")
     with gr.Group():
       msg = gr.Textbox(label = 'Search')
+      with gr.Accordion("Advanced Settings", open=False):
+        with gr.Row(equal_height = True):
+          llm_model = gr.Dropdown(choices = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'None'], 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 to sent as context")
+
     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(message):
+    def update_with_rag_md(message, llm_results_use = 5):
         rag_out = get_rag(message)
         md_text_updated = mark_text
-        for i in range(10):
+        for i in range(retrieve_results):
           rag_answer = rag_out[i]
           title = rag_answer['document_metadata']['title'].replace('\n','')
-          paper_title = f'''### [{title}](https://arxiv.org/abs/{rag_answer['document_id']})\n'''
+
+          score = round(rag_answer['score'], 2)
+          paper_title = f'''### **{score}** | [{title}](https://arxiv.org/abs/{rag_answer['document_id']})\n'''
           paper_abs = rag_answer['content']
-          md_text_updated += paper_title + paper_abs +  '\n---------------\n'+ '\n'
-        prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out[:llm_results]))
+          authors = rag_answer['document_metadata']['authors'].replace('\n','')
+          authors_formatted = f'*{authors}*' + ' \n\n'
+        
+          md_text_updated += paper_title + authors_formatted + paper_abs +  '\n---------------\n'+ '\n'
+        prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out[:llm_results_use]))
         return md_text_updated, prompt
 
-    def ask_llm(prompt):
+    def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
+       if llm_model_picked == 'None':
+          return gr.Textbox(visible = False)
+       client = InferenceClient(llm_model_picked)
        output = client.text_generation(prompt, **generate_kwargs,  stream=False, details=False, return_full_text=False)
        output = output.lstrip(' \n') if output.lstrip().startswith('\n') else output
        return gr.Textbox(output, visible = True)
 
-    msg.submit(update_with_rag_md, msg, [gr_md, input]).success(ask_llm, input, output_text)
+    msg.submit(update_with_rag_md, [msg, llm_results], [gr_md, input]).success(ask_llm, [input, llm_model], output_text)
 
 demo.launch(debug = True)