add image

app.py

@@ -12,6 +12,9 @@ import openai
 import google.generativeai as genai
 
 
+# Set up LLM APIs
+llm_api_options = ['gemini-pro', 'gemini-1.5-flash', 'gpt-3.5-turbo-1106']
+
 def query_gpt_model(
     prompt: str,
     llm: str = 'gpt-3.5-turbo-1106',
@@ -81,8 +84,8 @@ def query_model(
   else:
     raise ValueError('Unexpected model_name: ', model_name)
 
-# Load QuALITY dataset
 
+# Load QuALITY dataset
 _ONE2ONE_FIELDS = (
     'article',
     'article_id',
@@ -166,7 +169,6 @@ def quality_gutenberg_parser(raw_article):
   return ' '.join(lines)
 
 
-
 # ReadAgent (1) Episode Pagination
 prompt_pagination_template = """
 You are given a passage that is taken from a larger text (article, book, ...) and some numbered labels between the paragraphs in the passage.
@@ -250,8 +252,6 @@ def quality_pagination(example,
   text_output += f"\n\n[Pagination] Done with {len(pages)} pages"
   return pages, text_output
 
-# pages = quality_pagination(example)
-
 
 # ReadAgent (2) Memory Gisting
 prompt_shorten_template = """
@@ -284,8 +284,6 @@ def quality_gisting(example, pages, model_name, client=None, word_limit=600, sta
   text_output += f"\n\ncompression rate {round(100.0 - gist_word_count/word_count*100, 2)}% ({gist_word_count}/{word_count})"
   return output, text_output
 
-# example_with_gists = quality_gisting(example, pages)
-
 
 # ReadAgent (3) Look-Up
 prompt_lookup_template = """
@@ -405,6 +403,7 @@ def quality_parallel_lookup(example, model_name, client, verbose=True):
     return text_outputs
 
 
+# ReadAgent
 def query_model_with_quality(
       index: int,
       model_name: str = 'gemini-pro',
@@ -421,38 +420,52 @@ def query_model_with_quality(
      genai.configure(api_key=api_key)
 
    example = quality_dev[index]
+   article = f"[Title: {example['title']}]\n\n{example['article']}"
    pages, pagination = quality_pagination(example, model_name, client)
    print('Finish Pagination.')
    example_with_gists, gisting = quality_gisting(example, pages, model_name, client)
    print('Finish Gisting.')
    answers = quality_parallel_lookup(example_with_gists, model_name, client)
-   return prompt_pagination_template, pagination, prompt_shorten_template, gisting, prompt_lookup_template, '\n\n'.join(answers)
+   # return prompt_pagination_template, pagination, prompt_shorten_template, gisting, prompt_lookup_template, '\n\n'.join(answers)
+   return article, pagination, gisting, '\n\n'.join(answers)
 
 
-llm_api_options = ['gemini-pro', 'gemini-1.5-flash', 'gpt-3.5-turbo-1106']
-
 with gr.Blocks() as demo:
     gr.Markdown(
     """
     # A Human-Inspired Reading Agent with Gist Memory of Very Long Contexts
+
+    [[website]](https://read-agent.github.io/)
+    [[view on huggingface]](https://huggingface.co/spaces/ReadAgent/read-agent)
+    [[arXiv]](https://arxiv.org/abs/2402.09727)
+    [[OpenReview]](https://openreview.net/forum?id=OTmcsyEO5G)
+
+    ![teaser](/file=./asset/teaser.png)
+
+    The demo below showcases a version of the ReadAgent algorithm, which is nspired by how humans interactively read long documents.    
+    We implement ReadAgent as a simple prompting system that uses the advanced language capabilities of LLMs to (1) decide what content to store together in a memory episode (**Episode Pagination**), (2) compress those memory episodes into short episodic memories called gist memories (**Memory Gisting**), and (3) take actions to look up passages in the original text if ReadAgent needs to remind itself of relevant details to complete a task (**Parallel Lookup and QA**)
+    This demo can handle long-document reading comprehension tasks ([QuALITY](https://arxiv.org/abs/2112.08608); max 6,000 words) efficiently.
+
+    To get started, you can choose an index of QuALITY dataset.
+    This demo uses Gemini API or OpenAI API so it requires the corresponding API key.
     """)
-    with gr.Tab('ReadAgent (QuALITY)'):
-        llm_options = gr.Radio(llm_api_options, label="Backend LLM API", value='gemini-pro')
-        llm_api_key = gr.Textbox(
-          label="Paste your OpenAI API key (sk-...) or Gemini API key",
-          lines=1,
-          type="password",
-        )
-        with gr.Row():
-            with gr.Column():
-                index = gr.Dropdown(list(range(len(quality_dev))), value=13, label="QuALITY Index",)
-                button = gr.Button("Execute")
-                prompt_pagination = gr.Textbox(label="Episode Pagination Prompt Template", lines=5)
-                pagination_results = gr.Textbox(label="Episode Pagination", lines=20)
-                prompt_gisting = gr.Textbox(label="Memory Gisting Prompt Template", lines=5)
-                gisting_results = gr.Textbox(label="Memory Gisting", lines=20)
-                prompt_lookup = gr.Textbox(label="Parallel Lookup Prompt Template", lines=5)
-                lookup_qa_results = gr.Textbox(label="Parallel Lookup and QA", lines=20)
+    with gr.Row():
+        with gr.Column():
+            llm_options = gr.Radio(llm_api_options, label="Backend LLM API", value='gemini-pro')
+            llm_api_key = gr.Textbox(
+              label="Paste your OpenAI API key (sk-...) or Gemini API key",
+              lines=1,
+              type="password",
+            )
+            index = gr.Dropdown(list(range(len(quality_dev))), value=13, label="QuALITY Index")
+            button = gr.Button("Execute")
+            original_article = gr.Textbox(label="Original Article", lines=20)
+            # prompt_pagination = gr.Textbox(label="Episode Pagination Prompt Template", lines=5)
+            pagination_results = gr.Textbox(label="(1) Episode Pagination", lines=20)
+            # prompt_gisting = gr.Textbox(label="Memory Gisting Prompt Template", lines=5)
+            gisting_results = gr.Textbox(label="(2) Memory Gisting", lines=20)
+            # prompt_lookup = gr.Textbox(label="Parallel Lookup Prompt Template", lines=5)
+            lookup_qa_results = gr.Textbox(label="(3) Parallel Lookup and QA", lines=20)
 
     button.click(
         fn=query_model_with_quality,
@@ -462,12 +475,16 @@ with gr.Blocks() as demo:
            llm_api_key,
         ],
         outputs=[
-          prompt_pagination, pagination_results,
-          prompt_gisting, gisting_results,
-          prompt_lookup, lookup_qa_results,
+          # prompt_pagination, pagination_results,
+          # prompt_gisting, gisting_results,
+          # prompt_lookup, lookup_qa_results,
+          original_article,
+          pagination_results,
+          gisting_results,
+          lookup_qa_results,
         ]
     )
 
 
 if __name__ == '__main__':
-    demo.launch()
+    demo.launch(allowed_paths=['./asset/teaser.png'])