app.py

import streamlit as st
import pandas as pd
from pipelines.keyphrase_extraction_pipeline import KeyphraseExtractionPipeline
from pipelines.keyphrase_generation_pipeline import KeyphraseGenerationPipeline
import orjson

from annotated_text.util import get_annotated_html
import re
import string


@st.cache(allow_output_mutation=True, show_spinner=False)
def load_pipeline(chosen_model):
    if "keyphrase-extraction" in chosen_model:
        return KeyphraseExtractionPipeline(chosen_model)
    elif "keyphrase-generation" in chosen_model:
        return KeyphraseGenerationPipeline(chosen_model)


def extract_keyphrases():
    st.session_state.keyphrases = pipe(st.session_state.input_text)
    st.session_state.history[f"run_{st.session_state.current_run_id}"] = {
        "run_id": st.session_state.current_run_id,
        "model": st.session_state.chosen_model,
        "text": st.session_state.input_text,
        "keyphrases": st.session_state.keyphrases,
    }
    st.session_state.current_run_id += 1


def get_annotated_text(text, keyphrases):
    for keyphrase in keyphrases:
        text = re.sub(
            rf"({keyphrase})([^A-Za-z])",
            rf"$K:{keyphrases.index(keyphrase)}\2",
            text,
            flags=re.I,
            count=1,
        )

    result = []
    for i, word in enumerate(text.split(" ")):
        if "$K" in word and re.search(
            "(\d+)$", word.translate(str.maketrans("", "", string.punctuation))
        ):
            result.append(
                (
                    re.sub(
                        r"\$K:\d+",
                        keyphrases[
                            int(
                                re.search(
                                    "(\d+)$",
                                    word.translate(
                                        str.maketrans("", "", string.punctuation)
                                    ),
                                ).group(1)
                            )
                        ],
                        word,
                    ),
                    "KEY",
                    "#21c354",
                )
            )
        else:
            if i == len(st.session_state.input_text.split(" ")) - 1:
                result.append(f" {word}")
            elif i == 0:
                result.append(f"{word} ")
            else:
                result.append(f" {word} ")
    return result


def render_output(layout, runs, reverse=False, multi_select=False):
    runs = list(runs.values())[::-1] if reverse else list(runs.values())
    for run in runs:
        layout.markdown(f"**⚙️ Output run {run.get('run_id')}**")

        layout.markdown(f"**Model**: {run.get('model')}")
        result = get_annotated_text(run.get("text"), list(run.get("keyphrases")))

        layout.markdown(
            get_annotated_html(*result),
            unsafe_allow_html=True,
        )
        if "generation" in st.session_state.chosen_model:
            abstractive_keyphrases = [
                keyphrase
                for keyphrase in run.get("keyphrases")
                if keyphrase.lower() not in run.get("text").lower()
            ]
            layout.write(", ".join(abstractive_keyphrases))
        layout.markdown("---")


if "config" not in st.session_state:
    with open("config.json", "r") as f:
        content = f.read()
    st.session_state.config = orjson.loads(content)
    st.session_state.history = {}
    st.session_state.keyphrases = []
    st.session_state.current_run_id = 1
    st.session_state.chosen_model = st.session_state.config.get("models")[0]

if "select_rows" not in st.session_state:
    st.session_state.selected_rows = []

st.set_page_config(
    page_icon="🔑",
    page_title="Keyphrase extraction/generation with Transformers",
    layout="centered",
)

with open("css/style.css") as f:
    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

st.header("🔑 Keyphrase extraction/generation with Transformers")

description = """
Keyphrase extraction is a technique in text analysis where you extract the important keyphrases
from a text. Since this is a time-consuming process, Artificial Intelligence is used to automate it.
Currently, classical machine learning methods, that use statistics and linguistics, are widely used
for the extraction process. The fact that these methods have been widely used in the community has
the advantage that there are many easy-to-use libraries. Now with the recent innovations in
deep learning methods (such as recurrent neural networks and transformers, GANS, …),
keyphrase extraction can be improved. These new methods also focus on the semantics and
context of a document, which is quite an improvement.

This space gives you the ability to test around with some keyphrase extraction and generation models.
Keyphrase extraction models are transformers models fine-tuned as a token classification problem where
the tokens in a text are annotated as:

* B: Beginning of a keyphrase
* I: Inside a keyphrases
* O: Outside a keyhprase.

While keyphrase extraction can only extract keyphrases from a given text. Keyphrase generation models
work a bit differently. Here you use an encoder-decoder model like BART to generate keyphrases from a given text.
These models also have the ability to generate keyphrases, which are not present in the text 🤯.

Do you want to see some magic 🧙‍♂️? Try it out yourself! 👇
"""

st.write(description)

with st.form("keyphrase-extraction-form"):
    selectbox_container, _ = st.columns(2)

    st.session_state.chosen_model = selectbox_container.selectbox(
        "Choose your model:", st.session_state.config.get("models")
    )

    st.markdown(
        f"For more information about the chosen model, please be sure to check out the [🤗 Model Card](https://huggingface.co/DeDeckerThomas/{st.session_state.chosen_model})."
    )

    st.session_state.input_text = st.text_area(
        "✍ Input", st.session_state.config.get("example_text"), height=250
    ).replace("\n", " ")

    with st.spinner("Extracting keyphrases..."):
        pressed = st.form_submit_button("Extract")

if pressed:
    with st.spinner("Loading pipeline..."):
        pipe = load_pipeline(
            f"{st.session_state.config.get('model_author')}/{st.session_state.chosen_model}"
        )
    with st.spinner("Extracting keyphrases"):
        extract_keyphrases()

options = st.multiselect(
    "Specify runs you want to see",
    st.session_state.history.keys(),
    format_func=lambda run_id: f"Run {run_id.split('_')[1]}",
)

if len(st.session_state.history.keys()) > 0:
    if options:
        render_output(
            st,
            {key: st.session_state.history[key] for key in options},
        )
    else:
        render_output(st, st.session_state.history, reverse=True)