add initial visualize app

.gitignore

@@ -0,0 +1,2 @@
+.env
+__pycache__

app.py

@@ -0,0 +1,228 @@
+import sys
+
+import gradio as gr
+
+sys.path.append(".")
+sys.path.append("..")
+sys.path.append("../..")
+
+from cluster import cluster
+from extract import extract_endpoint
+from generate_answers import generate_relevant_chunks
+
+queries = [
+    "What is the size, shape, and energy (watt hour) or capacity (Amp hour) of battery discussed in the paper?",
+    "What specific mechanical testing methods were used to quantify strength?",
+    "What parameters they used to quantify the benefit of their individual design (mass saving, increased run time, etc.)?",
+    "What material chemistry combination (on the anode, cathode, separator, and electrolyte) was used in these papers?",
+    "What kind of end use application they targeted?",
+]
+MAX_CATEGORIES = 10
+
+
+def change_button(text):
+    if len(text) > 0:
+        return gr.Button(interactive=True)
+    else:
+        return gr.Button(interactive=False)
+
+
+def generate_category_btn(cluster_output):
+    unique_categories = set()
+    for item in cluster_output:
+        unique_categories.update(item["categories"])
+
+    update_show = [gr.Button(visible=True, value=w) for w in unique_categories]
+    update_hide = [
+        gr.Button(visible=False, value="")
+        for _ in range(MAX_CATEGORIES - len(unique_categories))
+    ]
+    return update_show + update_hide
+
+
+def get_query(btn):
+    return btn
+
+
+btn_list = []
+
+
+with gr.Blocks() as app:
+    gr.Markdown(
+        """
+    # Paper Query Clustering + Visualization
+    This app extracts text from papers and then searches for relevant excerpts based on a query. It then clusters and visualizes the relevant excerpts to find common themes across the papers.
+
+    ### Input
+    1. A group of research papers that you want to run the query on.
+    1. Query that you would like to know about these papers.
+
+    ### Output
+    Clustering and visualization of the relevant excerpts which answer the query across the papers.
+
+    # 1. Upload + Extract
+    First, upload the papers you want to analyze. Currently, we only support PDFs. Once they're uploaded, you can extract the text data from the papers.
+    """
+    )
+    file_upload = gr.Files()
+    extract_btn = gr.Button("Extract", interactive=False)
+    with gr.Tab(label="Table"):
+        extract_df = gr.Dataframe(
+            datatype="markdown", column_widths=[100, 400], wrap=True
+        )
+    with gr.Tab(label="JSON"):
+        extract_output = gr.JSON(label="Extract Output")
+
+    gr.Markdown(
+        """
+    ----------------
+    # 2. Extract Relevant Excerpts
+    Enter a query about these papers. This will search the papers to find the most relevant excerpts.
+    """
+    )
+
+    gr.Markdown(
+        """
+    ### Input
+    """
+    )
+    query = gr.Textbox(
+        label="Query", value=queries[1], lines=3, placeholder="Enter a query"
+    )
+    gr.Markdown(
+        """
+    You can also select some example queries below.
+    """
+    )
+    with gr.Row():
+        q0_btn = gr.Button(queries[0])
+        q1_btn = gr.Button(queries[1])
+        q2_btn = gr.Button(queries[2])
+        q3_btn = gr.Button(queries[3])
+        q4_btn = gr.Button(queries[4])
+    gr.Markdown(
+        """
+    ----
+    """
+    )
+    relevant_btn = gr.Button("Extract Excerpts", interactive=False)
+    gr.Markdown(
+        """
+    ### Output
+    """
+    )
+    with gr.Tab(label="Output Table"):
+        relevant_df = gr.Dataframe(
+            datatype="markdown", column_widths=[100, 100, 300], wrap=True
+        )
+    with gr.Tab(label="JSON"):
+        relevant_output = gr.JSON(label="Relevant Chunks Output")
+
+    gr.Markdown(
+        """
+    ----------------
+    # 3. Cluster & Visualize
+    Cluster the relevant excerpts to find common themes and visualize the results.
+    """
+    )
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(
+                """
+            ### Input
+            """
+            )
+            cluster_btn = gr.Button("Cluster", interactive=False)
+            cluster_output = gr.JSON(label="Cluster Output", visible=False)
+
+    gr.Markdown(
+        """
+    ### Visualization
+    """
+    )
+    visualize_output = gr.Plot()
+    with gr.Row():
+        for i in range(MAX_CATEGORIES):
+            btn = gr.Button(visible=False)
+            btn_list.append(btn)
+    with gr.Tab(label="By Paper"):
+        cluster_df = gr.Dataframe(
+            datatype="markdown", column_widths=[100, 100, 300], wrap=True
+        )
+
+    with gr.Tab(label="By Excerpt"):
+        cluster_granular_df = gr.Dataframe(
+            datatype="markdown", column_widths=[100, 100, 300], wrap=True
+        )
+
+    # Event handlers
+    file_upload.change(fn=change_button, inputs=[file_upload], outputs=[extract_btn])
+
+    extract_btn.click(
+        fn=extract_endpoint,
+        inputs=[file_upload],
+        outputs=[extract_output, extract_df],
+    )
+
+    extract_output.change(
+        fn=change_button,
+        inputs=[extract_output],
+        outputs=[relevant_btn],
+    )
+
+    q0_btn.click(
+        fn=get_query,
+        inputs=[q0_btn],
+        outputs=[query],
+    )
+
+    q1_btn.click(
+        fn=get_query,
+        inputs=[q1_btn],
+        outputs=[query],
+    )
+
+    q2_btn.click(
+        fn=get_query,
+        inputs=[q2_btn],
+        outputs=[query],
+    )
+
+    q3_btn.click(
+        fn=get_query,
+        inputs=[q3_btn],
+        outputs=[query],
+    )
+
+    q4_btn.click(
+        fn=get_query,
+        inputs=[q4_btn],
+        outputs=[query],
+    )
+
+    relevant_btn.click(
+        fn=generate_relevant_chunks,
+        inputs=[query, extract_output],
+        outputs=[relevant_output, relevant_df],
+        api_name="relevant_chunks",
+    )
+
+    relevant_output.change(
+        fn=change_button, inputs=[relevant_output], outputs=[cluster_btn]
+    )
+
+    cluster_btn.click(
+        fn=cluster,
+        inputs=[query, relevant_output],
+        outputs=[cluster_output, cluster_df, visualize_output, cluster_granular_df],
+        api_name="cluster",
+    )
+
+    cluster_output.change(
+        fn=generate_category_btn,
+        inputs=[cluster_output],
+        outputs=btn_list,
+    )
+
+if __name__ == "__main__":
+    app.launch()

cluster.py

@@ -0,0 +1,159 @@
+import pprint
+
+import pandas as pd
+from uniflow.flow.client import TransformClient
+from uniflow.flow.config import TransformOpenAIConfig
+from uniflow.op.prompt import Context
+
+from helpers import compare_strings_ignore_non_string
+from visualize_upload import visualize
+
+DEBUG = False
+
+
+def cluster(query, answers_data):
+    answers = []
+    for answer in answers_data:
+        answers.extend(answer["answer"])
+
+    data = [Context(context=query, excerpts=answers)]
+
+    instruction = """
+# Task: I am a researcher with a query about research papers. I have a list of excerpts from those papers. I need you to cluster each of these excerpts into a category based on the query.
+## Input:
+1. context: A brief query/context
+2. excerpts: An list of excerpts from research papers.
+## Evaluation Steps:
+### Step 1
+Go through each excerpt. For each excerpt, if there is an answer to the context/query that's not already captured by a category, create a category and add it to your category list. If the context has the word 'specific', make the category as specific as the excerpt. Repeat this process for each excerpt. The categories should be mutually exclusive.
+### Step 2
+Once you've gone through all the excerpts and you have a list of categories, go through the excerpts a second time, and this time assign each excerpt to a category. A single excerpt can be assigned to multiple categories. If there is no information relevant to any of the categories, please categorize the excerpt as "None".
+## Response Format: Your response should only include two fields below:
+1. categories: A list of all the generated categories. This is the output of Step 1 above.
+2. clusters: An object, with each category as a key, and a list of all the excerpts as strings that fall into that category as the value. This is the output of Step 2 above.
+"""
+
+    few_shot_examples = [
+        # Context(
+        #     context="Which types of batteries are discussed?",
+        #     excerpts=[
+        #         "This investigation will shed lights on the tuneable chemical environments of transition-metal oxides for advanced cathode materials and promote the development of sodium-ion batteries.",
+        #         "Bi2Se3 was studied as a novel sodium-ion battery anode material because of its high theoretical capacity and high intrinsic conductivity.",
+        #         "Magnesium-ion batteries (MIBs) are considered strong candidates for next-generation energy-storage systems owing to their high theoretical capacity, divalent nature and the natural abundancy of magnesium (Mg) resources on Earth.",
+        #         "Magnesium-ion batteries (MIBs) have great potential in large-scale energy storage field with high capacity, excellent safety, and low cost.",
+        #     ],
+        #     categories=["Sodium-ion battery", "Magnesium-ion batteries"],
+        #     clusters={
+        #         "Sodium-ion battery": [
+        #             "This investigation will shed lights on the tuneable chemical environments of transition-metal oxides for advanced cathode materials and promote the development of sodium-ion batteries.",
+        #             "Bi2Se3 was studied as a novel sodium-ion battery anode material because of its high theoretical capacity and high intrinsic conductivity.",
+        #         ],
+        #         "Magnesium-ion batteries": [
+        #             "Magnesium-ion batteries (MIBs) are considered strong candidates for next-generation energy-storage systems owing to their high theoretical capacity, divalent nature and the natural abundancy of magnesium (Mg) resources on Earth.",
+        #             "Magnesium-ion batteries (MIBs) have great potential in large-scale energy storage field with high capacity, excellent safety, and low cost.",
+        #         ],
+        #     },
+        # ),
+        # Context(
+        #     context="Which 3D printing materials are discussed?",
+        #     excerpts=[
+        #         "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",
+        #         "To this end, this work designs a novel 3D printing phase change aggregate to prepare concrete with prominent thermal capacity and ductility.",
+        #         "In this study, 15 commercial pure titanium samples are processed under different conditions, and the 3D pore structures are characterized by X-ray tomography",
+        #         "In this study, a support-less ceramic printing (SLCP) process using a hydrogel bath was developed to facilitate the manufacture of complex bone substitutes.",
+        #     ],
+        #     categories=[
+        #         "metals",
+        #         "polymer composites",
+        #         "ceramics",
+        #         "concrete",
+        #         "phase change aggregate",
+        #     ],
+        #     clusters={
+        #         "metals": [
+        #             "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",
+        #             "In this study, 15 commercial pure titanium samples are processed under different conditions, and the 3D pore structures are characterized by X-ray tomography",
+        #         ],
+        #         "polymer composites": [
+        #             "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented"
+        #         ],
+        #         "ceramics": [
+        #             "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",
+        #             "In this study, a support-less ceramic printing (SLCP) process using a hydrogel bath was developed to facilitate the manufacture of complex bone substitutes.",
+        #         ],
+        #         "concrete": [
+        #             "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",
+        #             "To this end, this work designs a novel 3D printing phase change aggregate to prepare concrete with prominent thermal capacity and ductility.",
+        #         ],
+        #         "phase change aggregate": [
+        #             "To this end, this work designs a novel 3D printing phase change aggregate to prepare concrete with prominent thermal capacity and ductility."
+        #         ],
+        #     },
+        # ),
+    ]
+
+    num_thread_batch_size = 1
+
+    config = TransformOpenAIConfig()
+    config.prompt_template.instruction = instruction
+    config.prompt_template.few_shot_prompt = few_shot_examples
+    config.model_config.model_name = "gpt-4-1106-preview"
+    config.model_config.response_format = {"type": "json_object"}
+    config.model_config.num_call = 1
+    config.model_config.temperature = 0.0
+    config.model_config.num_thread = num_thread_batch_size
+    config.model_config.batch_size = num_thread_batch_size
+
+    cluster_client = TransformClient(config)
+
+    output = cluster_client.run(data)
+    if DEBUG:
+        pprint.pprint(output)
+    output_clusters = answers_data
+    clusters = output[0]["output"][0]["response"][0]["clusters"]
+    output_answer_category = []
+
+    for idx, paper in enumerate(answers_data):
+        # Initialize an empty list to store the categories for each answer
+        categories_per_answer = []
+
+        # Iterate over each answer
+        for ans in paper["answer"]:
+            categories = []
+            # Iterate over each category in clusters
+            for category, texts in clusters.items():
+                # Check if the answer is in any of the texts related to the category
+                if any(compare_strings_ignore_non_string(ans, text) for text in texts):
+                    if category not in categories_per_answer:
+                        categories.append(category)
+                    output_answer_category.append(
+                        {"paper": paper["paper"], "answer": ans, "category": category}
+                    )
+            if len(categories) == 0:
+                categories.append("None")
+            categories_per_answer.extend(categories)
+
+        output_clusters[idx]["categories"] = categories_per_answer
+    for output_cluster in output_clusters:
+        if len(output_cluster["categories"]) == 0:
+            output_cluster["categories"].append("None")
+    df = create_category_df(output_clusters, answers_data)
+    output_answer_category_df = pd.DataFrame(output_answer_category)
+    visualize_output = visualize(output_clusters)
+
+    return [output_clusters, df, visualize_output, output_answer_category_df]
+
+
+def create_category_df(cluster_output, answers_data):
+    pd_data = {
+        "Paper": [],
+        "Excerpts": [],
+        "Categories": [],
+    }
+    for i, paper in enumerate(cluster_output):
+        pd_data["Paper"].append(paper["paper"])
+        pd_data["Excerpts"].append(", ".join(answers_data[i]["answer"]))
+        pd_data["Categories"].append(", ".join(paper["categories"]))
+
+    df = pd.DataFrame(pd_data)
+    return df

extract.py

@@ -0,0 +1,57 @@
+import re
+
+import pandas as pd
+from dotenv import load_dotenv
+from llama_index.core import SimpleDirectoryReader
+from llama_parse import LlamaParse
+
+load_dotenv()
+MIN_PARAGRAPH_LENGTH = 50
+
+
+def extract_paragraphs(markdown_text):
+    """
+    Extract paragraphs from a markdown text.
+    """
+    # Split the text into paragraphs using regex
+    paragraphs = re.split(r"\n\n+", markdown_text)
+    # Remove leading and trailing whitespaces from each paragraph
+    paragraphs = [p.strip() for p in paragraphs if p.strip()]
+    paragraphs = [
+        p
+        for p in paragraphs
+        if len(p) >= MIN_PARAGRAPH_LENGTH and not p.startswith("#")
+    ]
+    print(f"created {len(paragraphs)} paragraphs\n", paragraphs)
+
+    return paragraphs
+
+
+def extract_endpoint(file_paths):
+    """
+    Extract PDFs using LlamaParse.
+    """
+
+    # set up parser
+    parser = LlamaParse(result_type="markdown")  # "markdown" and "text" are available
+
+    # use SimpleDirectoryReader to parse our file
+    file_extractor = {".pdf": parser}
+    documents = SimpleDirectoryReader(
+        input_files=file_paths, file_extractor=file_extractor
+    ).load_data()
+
+    extracted_data = []
+
+    for doc in documents:
+        print(doc.text[:500])
+        paragraphs = extract_paragraphs(doc.text)
+        data = {
+            "paper": doc.metadata["file_name"],
+            "chunks": paragraphs,
+        }
+        extracted_data.append(data)
+
+    df = pd.DataFrame(extracted_data)
+
+    return [extracted_data, df]

generate_answers.py

@@ -0,0 +1,76 @@
+import gradio as gr
+import pandas as pd
+from uniflow.flow.client import TransformClient
+from uniflow.flow.config import TransformOpenAIConfig
+from uniflow.op.prompt import Context
+
+DEBUG = False
+
+
+def generate_relevant_chunks(query, input_data, progress=gr.Progress()):
+    data_list = []
+    for paper in input_data:  # progress.tqdm(input_data, desc="Papers"):
+        data = [Context(context=query, paragraph=p) for p in paper["chunks"]]
+        data_list.append({"paper": paper["paper"], "data": data})
+
+    instruction = """
+# Task: I am a researcher trying to understand information across several research papers. You are to determine which of the chunks most directly contains information related to the query.
+## Input:
+1. context: A brief query or description of the information I am looking for.
+2. paragraph: An paragraph from a research paper.
+## Evaluation Criteria: You should pick which sentence(s) contains directly relevant information to the context. The best answer is the sentences that most directly answer or contain the information specific to the context. If there are no such sentences, you should answer with ["None"].
+## Response Format: Your response should only include two fields below:
+1. explanation: Reasoning behind your judgment, explaining why the answer is appropriate or not.
+2. answer: The best sentence(s) that meet the Evaluation Criteria as a list of strings. This should be ["None"] if no sentence answers the query. At most, include 3 sentences.
+"""
+
+    few_shot_examples = []
+
+    num_thread_batch_size = 16
+
+    config = TransformOpenAIConfig()
+    config.prompt_template.instruction = instruction
+    config.prompt_template.few_shot_prompt = few_shot_examples
+    config.model_config.model_name = "gpt-4-1106-preview"
+    config.model_config.response_format = {"type": "json_object"}
+    config.model_config.num_call = 1
+    config.model_config.temperature = 0.0
+    config.model_config.num_thread = num_thread_batch_size
+    config.model_config.batch_size = num_thread_batch_size
+
+    client = TransformClient(config)
+
+    output = []
+
+    for paper in data_list:
+        init_output = client.run(paper["data"])
+        combined_output = init_output[0]
+        combined_output["output"][0]["response"][0]["explanation"] = [
+            combined_output["output"][0]["response"][0]["explanation"]
+        ]
+        if DEBUG:
+            print(combined_output)
+        for item in init_output[1:]:
+            combined_output["output"][0]["response"][0]["answer"].extend(
+                item["output"][0]["response"][0]["answer"]
+            )
+            combined_output["output"][0]["response"][0]["explanation"].append(
+                item["output"][0]["response"][0]["explanation"]
+            )
+        output.append(combined_output)
+
+    output_answers = []
+
+    for idx, o in enumerate(output):
+        filtered_answers = [
+            item for item in o["output"][0]["response"][0]["answer"] if item != "None"
+        ]
+        if len(filtered_answers) == 0:
+            filtered_answers = ["None"]
+        output_answers.append(
+            {"paper": input_data[idx]["paper"], "answer": filtered_answers}
+        )
+
+    df = pd.DataFrame(output_answers)
+
+    return [output_answers, df]

helpers.py

@@ -0,0 +1,20 @@
+import unicodedata
+
+
+def _remove_non_string_characters(string):
+    symbols_to_remove = ["Δ"]
+    return "".join(
+        char
+        for char in string
+        if unicodedata.category(char)[0] in {"L", "N", "P", "Z"}
+        and char not in symbols_to_remove
+    )
+
+
+def compare_strings_ignore_non_string(string1, string2):
+    string1 = _remove_non_string_characters(string1)
+    string2 = _remove_non_string_characters(string2)
+    if string1 != string2 and string1[0:20] == string2[0:20]:
+        print(f"String1: {string1}")
+        print(f"String2: {string2}")
+    return string1 == string2

requirements.txt

@@ -0,0 +1,8 @@
+uniflow==0.0.25
+python-dotenv==1.0.1
+gradio==4.19.2
+llama-index==0.10.19
+llama-parse==0.3.9
+rapidfuzz==3.6.2
+dataclasses-json==0.6.4
+plotly==5.20.0

visualize_upload.py

@@ -0,0 +1,51 @@
+import pandas as pd
+import plotly.graph_objs as go
+
+
+def visualize(cluster_data):
+
+    # Create a DataFrame with counts of each category
+    category_counts = {}
+    for paper in cluster_data:
+        categories = paper["categories"]
+        for category in categories:
+            if category in category_counts:
+                category_counts[category] += 1
+            else:
+                category_counts[category] = 1
+
+    category_df = pd.DataFrame(
+        {
+            "Category": list(category_counts.keys()),
+            "Count": list(category_counts.values()),
+        }
+    )
+
+    # Sort the DataFrame by count in descending order
+    category_df = category_df.sort_values(by="Count", ascending=False)
+
+    # Create hover text containing the count and titles of all papers for each category
+    hover_text = []
+    for category in category_df["Category"]:
+        titles = []
+        for paper in cluster_data:
+            if category in paper["categories"]:
+                titles.append(f'<a href="https://plot.ly/">{paper["paper"]}</a>')
+        hover_text.append(f'<br>Papers:<br>{"<br>".join(titles)}')
+
+    # Create Plotly Bar chart
+    fig = go.Figure(
+        data=[
+            go.Bar(
+                x=category_df["Category"],
+                y=category_df["Count"],
+                hovertext=hover_text,
+                marker=dict(color="brown"),
+            )
+        ]
+    )
+
+    # Update layout
+    fig.update_layout(title="Categories", xaxis_title="Category", yaxis_title="Count")
+
+    return fig