Added download, streaming and initial placeholder

app.py

@@ -5,6 +5,7 @@ from ragatouille import RAGPretrainedModel
 from huggingface_hub import InferenceClient
 import re
 from datetime import datetime
+import json
 
 retrieve_results = 10 
 
@@ -16,15 +17,16 @@ generate_kwargs = dict(
     )
 
 RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
+
 try:
   gr.Info("Setting up retriever, please wait...")
-  _ = RAG.search("what is Mistral?", k = 1)
+  rag_initial_output = RAG.search("what is Mistral?", k = 1)
   gr.Info("Retriever working successfully!")
 except:
   gr.Warning("Retriever not working!")
 
 mark_text = '# 🔍 Search Results\n'
-header_text = "# ArXiv RAG\n"
+header_text = "# ArXivCS RAG \n"
 try:
   with open("README.md", "r") as f:
       mdfile = f.read()
@@ -36,6 +38,12 @@ try:
 except:
   pass
 
+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']
+
+
 def rag_cleaner(inp):
     rank = inp['rank']
     title = inp['document_metadata']['title']
@@ -59,15 +67,15 @@ def get_rag(message):
 with gr.Blocks(theme = gr.themes.Soft()) as demo:
     header = gr.Markdown(header_text)
     with gr.Group():
-      msg = gr.Textbox(label = 'Search')
+      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 = ['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)
+    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)
+    gr_md = gr.Markdown(mark_text + md_text_initial)
 
     def update_with_rag_md(message, llm_results_use = 5):
         rag_out = get_rag(message)
@@ -76,8 +84,9 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
           rag_answer = rag_out[i]
           title = rag_answer['document_metadata']['title'].replace('\n','')
 
-          score = round(rag_answer['score'], 2)
-          paper_title = f'''### **{score}** | [{title}](https://arxiv.org/abs/{rag_answer['document_id']})\n'''
+          #score = round(rag_answer['score'], 2)
+          date = rag_answer['document_metadata']['_time']
+          paper_title = f'''### {date} | [{title}](https://arxiv.org/abs/{rag_answer['document_id']}) | [⬇️](https://arxiv.org/pdf/{rag_answer['document_id']})\n'''
           paper_abs = rag_answer['content']
           authors = rag_answer['document_metadata']['authors'].replace('\n','')
           authors_formatted = f'*{authors}*' + ' \n\n'
@@ -90,9 +99,16 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
        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)
+       #output = client.text_generation(prompt, **generate_kwargs,  stream=False, details=False, return_full_text=False)
+       stream = client.text_generation(prompt, **generate_kwargs,  stream=True, details=True, return_full_text=False)
+       #output = output.lstrip(' \n') if output.lstrip().startswith('\n') else output
+       output = ""
+
+       for response in stream:
+          output += response.token.text
+          yield output
+       return output
+       #return gr.Textbox(output, visible = True)
 
     msg.submit(update_with_rag_md, [msg, llm_results], [gr_md, input]).success(ask_llm, [input, llm_model], output_text)