From 0b7b5839a78d0e814c7293c081c0cb9fa7093190 Mon Sep 17 00:00:00 2001
From: Hironsan <hiroki.nakayama.py@gmail.com>
Date: Mon, 27 Mar 2023 08:52:36 +0900
Subject: [PATCH] Add a question answering page

---
 doccano_mini/examples.py                    | 14 +++++++++++
 doccano_mini/pages/02_Question_Answering.py | 28 +++++++++++++++++++++
 doccano_mini/prompts.py                     | 22 ++++++++++++++++
 3 files changed, 64 insertions(+)
 create mode 100644 doccano_mini/pages/02_Question_Answering.py

diff --git a/doccano_mini/examples.py b/doccano_mini/examples.py
index 5d78531..5051981 100644
--- a/doccano_mini/examples.py
+++ b/doccano_mini/examples.py
@@ -13,6 +13,20 @@ def make_classification_example() -> pd.DataFrame:
     return df
 
 
+def make_question_answering_example() -> pd.DataFrame:
+    df = pd.DataFrame(
+        [
+            {
+                "context": "Google was founded by computer scientists Larry Page and Sergey Brin.",
+                "question": "Who founded Google?",
+                "answer": "Larry Page and Sergey Brin",
+            },
+        ],
+        columns=["context", "question", "answer"],
+    )
+    return df
+
+
 def make_task_free_example() -> pd.DataFrame:
     df = pd.DataFrame([{"Column 1": "", "Column 2": ""}], columns=["Column 1", "Column 2"])
     return df
diff --git a/doccano_mini/pages/02_Question_Answering.py b/doccano_mini/pages/02_Question_Answering.py
new file mode 100644
index 0000000..b8a92df
--- /dev/null
+++ b/doccano_mini/pages/02_Question_Answering.py
@@ -0,0 +1,28 @@
+from typing import Dict, List
+
+import streamlit as st
+
+from doccano_mini.examples import make_question_answering_example
+from doccano_mini.layout import BasePage
+from doccano_mini.prompts import make_question_answering_prompt
+
+
+class QuestionAnsweringPage(BasePage):
+    def make_examples(self, columns: List[str]):
+        df = make_question_answering_example()
+        edited_df = st.experimental_data_editor(df, num_rows="dynamic", width=1000)
+        examples = edited_df.to_dict(orient="records")
+        return examples
+
+    def make_prompt(self, examples: List[Dict]):
+        return make_question_answering_prompt(examples)
+
+    def prepare_inputs(self, columns: List[str]):
+        return {
+            "context": st.text_area(label="Context.", value="", height=300),
+            "question": st.text_area(label="Question.", value="", height=100),
+        }
+
+
+page = QuestionAnsweringPage(title="Question Answering")
+page.render()
diff --git a/doccano_mini/prompts.py b/doccano_mini/prompts.py
index af3b361..26001fc 100644
--- a/doccano_mini/prompts.py
+++ b/doccano_mini/prompts.py
@@ -23,6 +23,28 @@ def make_classification_prompt(examples: List[dict]) -> FewShotPromptTemplate:
     return prompt
 
 
+def make_question_answering_prompt(examples: List[dict]) -> FewShotPromptTemplate:
+    task_instruction = (
+        "You are a highly intelligent question answering bot. "
+        "You take context and question as input and return the answer from the context. "
+        "Retain as much information as needed to answer the question at a later time. "
+        "If you don't know the answer, you should return N/A."
+    )
+
+    example_prompt = PromptTemplate(
+        input_variables=["context", "question", "answer"],
+        template="context: {context}\nquestion: {question}\nanswer: {answer}",
+    )
+    prompt = FewShotPromptTemplate(
+        examples=examples,
+        example_prompt=example_prompt,
+        prefix=task_instruction,
+        suffix="context: {context}\nquestion: {question}",
+        input_variables=["context", "question"],
+    )
+    return prompt
+
+
 def make_task_free_prompt(examples: List[dict]) -> FewShotPromptTemplate:
     columns = list(examples[0])