app.py

import json

import streamlit as st
from annotated_text import annotated_text
from cherche import compose, qa, rank, retrieve, summary
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from transformers import pipeline


@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)
def loading_pipelines():
    """Create three pipelines dedicated to neural research. The first one is dedicated to game
    retrieval. The second is dedicated to the question answering task. The third is dedicated to
    the summarization task. Save pipelines as pickle file.

    >>> search = (
    ...    tfidf(on = "game") + ranker(on = "game") | tfidf(on = ["game", "summary"]) +
    ...    ranker(on = ["game", "summary"]) + documents
    ... )

    """
    # Load documents
    with open("games.json", "r") as documents_file:
        documents = json.load(documents_file)

    # A first retriever dedicated to title
    retriever_title = retrieve.TfIdf(
        key="id",
        on=["game"],
        documents=documents,
        tfidf=TfidfVectorizer(
            lowercase=True,
            min_df=1,
            max_df=0.9,
            ngram_range=(3, 7),
            analyzer="char",
        ),
        k=30,
    )

    # A second retriever dedicated to title and also summary of games.
    retriever_title_summary = retrieve.TfIdf(
        key="id",
        on=["game", "summary"],
        documents=documents,
        tfidf=TfidfVectorizer(
            lowercase=True,
            min_df=1,
            max_df=0.9,
            ngram_range=(3, 7),
            analyzer="char",
        ),
        k=30,
    )

    # Load our encoder to re-rank retrievers documents.
    encoder = SentenceTransformer("sentence-transformers/all-mpnet-base-v2").encode

    # A ranker dedicated to title
    ranker_title = rank.Encoder(
        key="id",
        on=["game"],
        encoder=encoder,
        k=5,
        path="games_title.pkl",
    )

    # A ranker dedicated to title and summary
    ranker_title_summary = rank.Encoder(
        key="id",
        on=["game", "summary"],
        encoder=encoder,
        k=5,
        path="games_summary.pkl",
    )

    # Pipeline creation
    search = (
        (retriever_title + ranker_title) | (retriever_title_summary + ranker_title_summary)
    ) + documents

    # Index
    search.add(documents)
    return search


@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)
def write_search(query):
    return search(query)[:5]


@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)
def loading_summarization_pipeline():
    summarizer = summary.Summary(
        model=pipeline(
            "summarization",
            model="sshleifer/distilbart-cnn-12-6",
            tokenizer="sshleifer/distilbart-cnn-12-6",
            framework="pt",
        ),
        on=["game", "summary"],
        max_length=50,
    )

    search_summarize = search + summarizer
    return search_summarize


@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)
def write_search_summarize(query_summarize):
    return search_summarize(query_summarize)


@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)
def loading_qa_pipeline():
    question_answering = qa.QA(
        model=pipeline(
            "question-answering",
            model="deepset/roberta-base-squad2",
            tokenizer="deepset/roberta-base-squad2",
        ),
        k=3,
        on="summary",
    )
    search_qa = search + question_answering
    return search_qa


@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)
def write_search_qa(query_qa):
    return search_qa(query_qa)


if __name__ == "__main__":

    st.markdown("# 🕹 Cherche")

    st.markdown(
        "[Cherche](https://github.com/raphaelsty/cherche) (search in French) allows you to create a \
        neural search pipeline using retrievers and pre-trained language models as rankers. Cherche's main strength is its ability to build diverse and end-to-end pipelines."
    )

    st.image("explain.png")

    st.markdown(
        "Here is a demo of neural search for video games using a sample of reviews made by [Metacritic](https://www.metacritic.com). \
        Starting the app may take a while if the models are not stored in cache."
    )

    # Will be slow the first time, you will need to compute embeddings.
    search = loading_pipelines()

    st.markdown("## 👾 Neural search")

    st.markdown(
        '```search = (tfidf(on = "title") + ranker(on = "title") | tfidf(on = ["title", "summary"]) + ranker(on = ["game", "summary"]) + documents)```'
    )

    query = st.text_input(
        "games",
        value="super smash bros",
        max_chars=None,
        key=None,
        type="default",
        help=None,
        autocomplete=None,
        on_change=None,
        args=None,
        kwargs=None,
    )

    if query:

        for document in write_search(query):
            if document["rate"] < 10:
                document["rate"] *= 10

            st.markdown(f"### {document['game']}")
            st.markdown(f"Metacritic Rating: {document['rate']}")

            col_1, col_2 = st.columns([1, 5])
            with col_1:
                st.image(document["image"], width=100)
            with col_2:
                st.write(f"{document['summary'][:430]}...")

    st.markdown("## 🎲 Summarization")

    st.markdown(
        '```search = (tfidf(on = "title") + ranker(on = "title") | tfidf(on = ["title", "summary"]) + ranker(on = ["game", "summary"]) + documents + summarization(on = "summary"))```'
    )

    st.markdown(
        "Let's create a summay but it may take few seconds. Summarization models are not that fast using CPU. Also it may take time to load the summarization model if it's not in cache yet.."
    )

    query_summarize = st.text_input(
        "summarization",
        value="super smash bros",
        max_chars=None,
        key=None,
        type="default",
        help=None,
        autocomplete=None,
        on_change=None,
        args=None,
        kwargs=None,
    )

    if query_summarize:
        search_summarize = loading_summarization_pipeline()
        st.write(f"**{write_search_summarize(query_summarize)}**")

    st.markdown("## 🎮 Question answering")

    st.markdown(
        '```search = (tfidf(on = "title") + ranker(on = "title") | tfidf(on = ["title", "summary"]) + ranker(on = ["game", "summary"]) + documents + question_answering(on = "summary"))```'
    )

    st.markdown(
        "It may take few seconds. Question answering models are not that fast using CPU. Also it may take time to load the question answering model if it's not in cache yet."
    )

    query_qa = st.text_input(
        "question",
        value="What is the purpose of playing Super Smash Bros?",
        max_chars=None,
        key=None,
        type="default",
        help=None,
        autocomplete=None,
        on_change=None,
        args=None,
        kwargs=None,
    )

    if query_qa:

        search_qa = loading_qa_pipeline()
        for document_qa in write_search_qa(query_qa):

            st.markdown(f"### {document_qa['game']}")
            st.markdown(f"Metacritic Rating: {document_qa['rate']}")

            col_1, col_2 = st.columns([1, 5])
            with col_1:
                st.image(document_qa["image"], width=100)
            with col_2:

                annotations = document_qa["summary"].split(document_qa["answer"])

                if document_qa["start"] == 0:
                    annotated_text(
                        (
                            document_qa["answer"],
                            f"answer {round(document_qa['qa_score'], 2)}",
                            "#8ef",
                        ),
                        " ",
                        " ".join(annotations[1:]),
                    )

                elif document_qa["end"] == len(document_qa["summary"]):
                    annotated_text(
                        " ".join(annotations[:-1]),
                        (
                            document_qa["answer"],
                            f"answer {round(document_qa['qa_score'], 2)}",
                            "#8ef",
                        ),
                    )

                else:
                    annotated_text(
                        annotations[0],
                        (
                            document_qa["answer"],
                            f"answer {round(document_qa['qa_score'], 2)}",
                            "#8ef",
                        ),
                        annotations[1],
                    )