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], )