app.py

"""
This application demo shows how to extract structured information using LLMs
and transfer it as metadata in Kadi.
"""

import os
import json
import gradio as gr
import groq
from difflib import Differ
from json2kadi import my_json_to_kadi
from kadi_apy.lib.conversion import json_to_kadi

# Set api key of Groq
api_key = os.getenv("GROQ_API")

# Examples
example_1 = (
    """John B. Goodenough (1922–2023) was a renowned American physicist and materials scientist,
best known for his pioneering work in developing the lithium-ion battery. He earned a Ph.D. in physics from
the University of Chicago in 1952. Throughout his career, Goodenough worked at prominent institutions and companies,
including the Massachusetts Institute of Technology (MIT) Lincoln Laboratory, where he helped develop
random-access memory (RAM), and later at the University of Oxford and the University of Texas at Austin.
He received the Nobel Prize in Chemistry in 2019 at the age of 97.""",
    """{
    "Name": "",
    "Birthday": "",
    "Educations": [
        {
            "School": "",
            "Date": ""
        }
    ],
    "Experiences": [
        {
            "Company": "",
            "Date": ""
        }
    ]
}""",
)

example_2 = (
    """Argyrodite-type Lithium Thiophosphate (Li₁₀GeP₂S₁₂) is a sulfide-based solid-state electrolyte that exhibits an impressive ionic conductivity of 10⁻³ S/cm, rivaling that of liquid electrolytes. Li₁₀GeP₂S₁₂ is known for its flexibility and ease of processing, which makes it highly adaptable to various solid-state battery architectures. Batteries using this electrolyte have shown a specific capacity of 180 mAh/g and retain 80% capacity after 700 cycles under ambient conditions. This material is seen as a key enabler for the development of high-performance, all-solid-state lithium batteries.""",
    """{
    "Material": {
        "Name": "",
        "Composition": "",
        "Type": "",
        "Properties": {
            "Ionic Conductivity": {"Value": "", "Unit": ""},
            "Chemical Stability": "",
            "Dendrite Formation Risk": "",
            "Operating Voltage": "",
            "Flexibility": "",
            "Processing": ""
        }
    },
    "Performance": {
        "Specific Capacity": {"Value": "", "Unit": ""},
        "Energy Density": {"Value": "", "Unit": ""},
        "Capacity Retention": "",
        "Operating Temperature": {"Value": "", "Unit": ""}
    },
    "Usage": {
        "Battery Type": "",
        "Benefits": []
    }
}""",
)

example_3 = (
    """In this experiment, LATP (Lithium Aluminum Titanium Phosphate) electrolyte was synthesized using a modified sol-gel method. Lithium acetate dihydrate (Li(C2H3O2)·2H2O), aluminum nitrate nonahydrate (Al(NO3)3·9H2O), and titanium isopropoxide (Ti[OCH(CH3)2]4) were used as precursors. Lithium acetate and aluminum nitrate were first dissolved in distilled water under constant stirring at room temperature. Titanium isopropoxide was then added dropwise to the solution. Phosphoric acid (H3PO4) was introduced slowly via a drip funnel to form a white gel, which was dried at room temperature for 24 hours. The dried gel was heat treated in two stages. Initially, it was heated to 400°C for 6 hours to remove volatile compounds and induce precursor formation. Subsequently, the material was heated to 900°C for 8 hours to complete the crystallization of LATP. The resultant powder was further processed using a planetary ball mill to enhance particle uniformity. To prepare pellets, the powders were pressed uniaxially and subsequently densified using cold isostatic pressing at 400 MPa. The green density of the pressed samples was approximately 62% relative density.""",
    """{
    "Material": {
        "Name": "",
        "Formula": "",
        "SynthesisMethod": "",
        "Precursors": [
            {"Name": "", "Formula": "", "Source": ""}
        ],
        "Solvent": {
            "Type": "",
            "Volume": {"Value": "", "Unit": ""}
        },
        "Mixing": {
            "Temperature": {"Value": "", "Unit": ""},
            "Duration": {"Value": "", "Unit": ""}
        },
        "Addition": {
            "Component": "",
            "Method": "",
            "Rate": {"Value": "", "Unit": ""}
        },
        "Drying": {
            "Temperature": {"Value": "", "Unit": ""},
            "Duration": {"Value": "", "Unit": ""}
        }
    },
    "HeatTreatment": [
        {
            "Temperature": {"Value": "", "Unit": ""},
            "Duration": {"Value": "", "Unit": ""},
            "Purpose": ""
        }
    ],
    "PostProcessing": {
        "Milling": "",
        "Pressing": {
            "Method": "",
            "Pressure": {"Value": "", "Unit": ""}
        }
    }
}
""",
)

example_4 = (
    """In this study, X-ray diffraction (XRD) was utilized to characterize the crystalline structure of the solid-state battery materials. A powdered sample of the synthesized electrolyte was prepared by grinding it into a fine powder using a mortar and pestle. The XRD measurements were conducted using a Bruker D8 Advance diffractometer equipped with a Cu Kα radiation source (λ = 1.5406 Å) operating at 40 kV and 40 mA. The diffraction data were collected over a 2θ range of 10° to 80° with a step size of 0.02° and a counting time of 1 second per step. The scan rate was set at 1° per minute. The obtained XRD patterns were analyzed to determine the phase purity and crystal structure of the electrolyte. Data analysis was performed using the Bruker EVA software, and phase identification was confirmed by comparing the experimental patterns to standard reference patterns from the International Centre for Diffraction Data (ICDD) database. The results provided insights into the phase composition, crystallite size, and lattice parameters of the material, which are crucial for understanding its performance in solid-state battery applications.""",
    """{
    "Material": {
        "Name": "",
        "Type": "",
        "Preparation": {
            "Method": "",
            "Details": ""
        },
        "Characterization": {
            "Technique": "",
            "Instrument": "",
            "Parameters": {
                "Radiation Source": {"Type": "", "Wavelength": {"Value": "", "Unit": ""}},
                "Voltage": {"Value": "", "Unit": ""},
                "Current": {"Value": "", "Unit": ""},
                "Scan Range": {"Start": {"Value": "", "Unit": ""}, "End": {"Value": "", "Unit": ""}},
                "Step Size": {"Value": "", "Unit": ""},
                "Counting Time": {"Value": "", "Unit": ""},
                "Scan Rate": {"Value": "", "Unit": ""}
            },
            "Data Analysis": {
                "Software": "",
                "Reference Database": ""
            }
        }
    }
}
""",
)


def generate_response(prompt):
    """
    Get response (structured json) from LLMs.
    """
    if not prompt:
        return "No transcription available. Please try speaking again."

    client = groq.Client(api_key=api_key)

    try:
        # Use Llama 3.1 70B powered by Groq for text generation
        completion = client.chat.completions.create(
            model="llama-3.1-70b-versatile",
            messages=[
                # {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ],
        )
        return completion.choices[0].message.content
    except Exception as e:
        return f"Error in response generation: {str(e)}"


def post_process_output(output):
    """Clean up output."""

    # 1. remove json mark
    output = output.replace("```", "")
    output = output.replace("null", '""')
    output = output.removeprefix("json")

    # remove trailing space etc.
    output = output.strip()
    return output


# Basic prompt for extraction
extract_info_prompt = """
You are an data scientist, extract information from text with given template in json format. Do not add any explanation.

Text:
{}

Template:
{}
"""


def extract_info(input_text, structure_template):
    """Extract structured output from text input."""

    # validate structure_template is json
    try:
        structure_template = json.dumps(json.loads(structure_template), indent=4)
    except:
        raise gr.Error("Error JSON schema in structure template.")

    combined_prompt = extract_info_prompt.format(input_text, structure_template)

    # try 3 times in case of json output error
    for _ in range(3):
        try:
            structured_output = generate_response(combined_prompt)
            structured_output = post_process_output(structured_output)
            structured_output = json.dumps(json.loads(structured_output), indent=4)
            break
        except Exception as e:
            print("Error in json format, retrying...")
            continue

    return structured_output


def diff_texts(text1, text2):
    """Compare two text inputs."""

    d = Differ()
    return [
        (token[2:], token[0] if token[0] != " " else None)
        for token in d.compare(text1, text2)
    ]


def transform_json_to_kadi_schema(input_json_str):
    """Tranform json into Kadi metadata schema."""

    input_json = json.loads(input_json_str)
    try:
        output_json = my_json_to_kadi(input_json)
    except Exception:
        # fallback to json_to_kadi from kadi_apy
        output_json = json_to_kadi(input_json)

    return json.dumps(output_json, indent=2)


# Baisc template for inferring json template
example_structure_template = """
{
    "Material": {
        "Name": "",
        "Type": "",
        "Properties": {
            "Ionic Conductivity": {"Value": "", "Unit": ""},
            "Chemical Stability": "",
            "Dendrite Formation Risk": "",
            "Operating Voltage": {"Value": "", "Unit": ""},
            "Flexibility": "",
            ""Processing": "",
        }
    },
    "Performance": {
        "Specific Capacity": {"Value": "", "Unit": ""},
        "Capacity Retention": {"Value": "", "Unit": ""},
        "Operating Temperature": {"Value": "", "Unit": ""}
    },
        "Usage": {
            "Battery Type": "",
            "Benefits": []
        }
    }
"""


# Infer template from text input based on exmaple template defined above
def suggest_template(input_text):
    """Infer structured template from text input."""

    if not input_text.strip():
        raise gr.Error("The input text should not be empty.")
    combined_prompt = f"""
        Extract and generalize a template from the provided text, following the structure of the given Example Templates.
        Replace specific values with placeholders. The output should be a template without concrete data. Maintain the format
        and style of the original Example Templates, and change items according to the text when necessary.

        Text:
        {input_text}

        Example Template:
        {example_structure_template}

        Output the result in json format. Do not give any explanation and do not give other information.
        """

    # try 3 times in case of json output error
    for _ in range(3):
        try:
            output = generate_response(combined_prompt)
            output = post_process_output(output)
            output = json.dumps(json.loads(output), indent=4)
            break
        except Exception as e:
            print("Error in json format, retrying...")
            continue
    return output


# Graio UI
with gr.Blocks() as demo:
    gr.Markdown(
        "### A simple web app to obtain structured output from text input using Large Language Models (LLMs)."
    )
    gr.Markdown(
        "⚠️ This is designed for research purposes, for best privacy protection, please avoid sharing sensitive or confidential data."
    )

    with gr.Tab("Text Extract 🗃️"):
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(
                    label="Text to extract",
                    lines=5,
                    placeholder="Enter your text here and click the button below to extract structured info.",
                )

                structure_template = gr.Textbox(
                    label="Structure template",
                    lines=3,
                    placeholder="Enter your structure template here.",
                )

                with gr.Row():
                    suggest_btn = gr.Button("Suggest template", scale=1)
                    submit_btn = gr.Button("Extract", variant="primary", scale=2)

            with gr.Column():
                output = gr.Textbox(label="Structured Output", show_copy_button=True)
                with gr.Accordion("Show Kadi-compatible output", open=False):
                    output_kadi = gr.Textbox(
                        label="Kadi compatible metadata output",
                        lines=5,
                        show_copy_button=True,
                    )

                    gr.Markdown()
                    gr.Markdown(
                        "Add metadata by copying and pasting in [Kadi](https://kadi.iam.kit.edu/) Record"
                    )
                    gr.Markdown("![](file/copy_to_kadi.png)")

        # Actions
        submit_btn.click(
            fn=extract_info, inputs=[text_input, structure_template], outputs=output
        )
        suggest_btn.click(
            fn=suggest_template, inputs=[text_input], outputs=structure_template
        )

        output.change(
            fn=transform_json_to_kadi_schema, inputs=[output], outputs=output_kadi
        )

        # Placeholder
        gr.Markdown()
        gr.Markdown()
        gr.Markdown()

        gr.Examples(
            examples=[
                [example_1[0], example_1[1]],
                [example_2[0], example_2[1]],
                [example_3[0], example_3[1]],
                [example_4[0], example_4[1]],
            ],
            examples_per_page=2,
            inputs=[text_input, structure_template],
        )

    with gr.Tab("About"):
        gr.Markdown(
            """
                This simple app is designed for academic purposes to explore the possibility of using Large Language Models (LLMs) to assist in research data management. It is being developed by a team from Karlsruhe, Germany, who are interested in using machine learning techniques to support various scientific topics and simplify the process of research data management. You may find more information about us [here](https://kadi.iam.kit.edu/#).
                
                **For best privacy protection, please avoid sharing sensitive or confidential data.**
                """
        )


if __name__ == "__main__":
    demo.launch(show_api=False, allowed_paths=["./"])