Merge pull request #9 from doccano/feature/task-free

Add the new task type Task Free

doccano_mini/app.py

@@ -4,14 +4,24 @@
 
 from doccano_mini.examples import select_example_df
 from doccano_mini.prompts import select_prompt_maker
-from doccano_mini.tasks import options
+from doccano_mini.tasks import TaskType, options
 
 CODE = """from langchain.chains import load_chain
 
 chain = load_chain("chain.yaml")
 chain.run("YOUR TEXT")"""
 
 
+# https://platform.openai.com/docs/models/gpt-3-5
+AVAILABLE_MODELS = (
+    "gpt-3.5-turbo",
+    "gpt-3.5-turbo-0301",
+    "text-davinci-003",
+    "text-davinci-002",
+    "code-davinci-002",
+)
+
+
 def display_download_button():
     st.header("Download LangChain's config")
     with open("config.yaml", "r", encoding="utf-8") as f:
@@ -22,14 +32,7 @@ def display_download_button():
         )
 
 
-def main():
-    st.set_page_config(page_title="doccano-mini", page_icon=":memo:")
-
-    st.title("doccano-mini")
-
-    st.header("Select your task")
-    task = st.selectbox("", options=options, label_visibility="collapsed")
-
+def task_classification(task: TaskType):
     st.header("Annotate your data")
     df = select_example_df(task)
     edited_df = st.experimental_data_editor(df, num_rows="dynamic", width=1000)
@@ -47,24 +50,16 @@ def main():
     col1, col2 = st.columns([3, 1])
     text = col1.text_area(label="Please enter your text.", value="", height=300)
 
-    # https://platform.openai.com/docs/models/gpt-3-5
-    available_models = (
-        "gpt-3.5-turbo",
-        "gpt-3.5-turbo-0301",
-        "text-davinci-003",
-        "text-davinci-002",
-        "code-davinci-002",
-    )
     # Use text-davinci-003 by default.
-    model_name = col2.selectbox("Model", available_models, index=2)
+    model_name = col2.selectbox("Model", AVAILABLE_MODELS, index=2)
     temperature = col2.slider("Temperature", min_value=0.0, max_value=1.0, value=0.7, step=0.01)
     top_p = col2.slider("Top-p", min_value=0.0, max_value=1.0, value=1.0, step=0.01)
 
     with st.expander("See your prompt"):
         st.markdown(f"```\n{prompt.format(input=text)}\n```")
 
     if st.button("Predict"):
-        llm = OpenAI(model_name=model_name, temperature=temperature, top_p=top_p)
+        llm = OpenAI(model_name=model_name, temperature=temperature, top_p=top_p)  # type:ignore
         chain = LLMChain(llm=llm, prompt=prompt)
         response = chain.run(text)
         label = response.split(":")[1]
@@ -73,6 +68,57 @@ def main():
         chain.save("config.yaml")
         display_download_button()
 
+
+def task_free(task: TaskType):
+    st.header("Annotate your data")
+
+    num_cols = st.number_input("Set the number of columns", min_value=2, max_value=10)
+    columns = [st.text_input(f"Column {i}:", value=f"column {i}") for i in range(1, int(num_cols) + 1)]
+
+    df = select_example_df(task)
+    df = df.reindex(columns, axis="columns", fill_value="")
+    edited_df = st.experimental_data_editor(df, num_rows="dynamic", width=1000)
+    examples = edited_df.to_dict(orient="records")
+
+    prompt = select_prompt_maker(task)(examples)
+
+    prompt.prefix = st.text_area(
+        label="Enter task instruction",
+        placeholder=f"Predict {columns[-1]} based on {', '.join(columns[:-1])}.",
+        height=200,
+    )
+
+    inputs = {column: st.text_input(f"Input for {column}:") for column in columns[:-1]}
+
+    st.markdown(f"Your prompt\n```\n{prompt.format(**inputs)}\n```")
+
+    # Use text-davinci-003 by default.
+    model_name = st.selectbox("Model", AVAILABLE_MODELS, index=2)
+    temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.7, step=0.01)
+    top_p = st.slider("Top-p", min_value=0.0, max_value=1.0, value=1.0, step=0.01)
+    if st.button("Predict"):
+        llm = OpenAI(model_name=model_name, temperature=temperature, top_p=top_p)  # type:ignore
+        chain = LLMChain(llm=llm, prompt=prompt)
+        response = chain.run(**inputs)
+        st.text(response)
+
+        chain.save("config.yaml")
+        display_download_button()
+
+
+def main():
+    st.set_page_config(page_title="doccano-mini", page_icon=":memo:")
+
+    st.title("doccano-mini")
+
+    st.header("Select your task")
+    task = st.selectbox("", options=options, label_visibility="collapsed")
+
+    if task == TaskType.TEXT_CLASSIFICATION.value:
+        task_classification(task)
+    else:
+        task_free(task)
+
     st.header("Usage")
     st.code(CODE)
 

doccano_mini/examples.py

@@ -14,4 +14,7 @@ def select_example_df(task: TaskType) -> pd.DataFrame:
             columns=["text", "label"],
         )
         return df
+    elif task == TaskType.TASK_FREE.value:
+        df = pd.DataFrame([{"Column 1": "", "Column 2": ""}], columns=["Column 1", "Column 2"])
+        return df
     raise ValueError(f"Task {task} is not supported.")

doccano_mini/prompts.py

@@ -24,7 +24,25 @@ def make_classification_prompt(examples: List[dict]) -> FewShotPromptTemplate:
     return prompt
 
 
+def make_task_free_prompt(examples: List[dict]) -> FewShotPromptTemplate:
+    columns = list(examples[0])
+
+    example_prompt = PromptTemplate(
+        input_variables=columns, template="\n".join([f"{column}: {{{column}}}" for column in columns])
+    )
+
+    prompt = FewShotPromptTemplate(
+        examples=examples,
+        example_prompt=example_prompt,
+        suffix="\n".join([f"{column}: {{{column}}}" for column in columns[:-1]]),
+        input_variables=columns[:-1],
+    )
+    return prompt
+
+
 def select_prompt_maker(task: TaskType):
     if task == TaskType.TEXT_CLASSIFICATION.value:
         return make_classification_prompt
+    elif task == TaskType.TASK_FREE.value:
+        return make_task_free_prompt
     raise ValueError(f"Task {task} is not supported.")

doccano_mini/tasks.py

@@ -3,6 +3,7 @@
 
 class TaskType(Enum):
     TEXT_CLASSIFICATION = "Text Classification"
+    TASK_FREE = "Task Free"
 
 
 options = [task_type.value for task_type in TaskType]