raphaelsty/games
4
1import json2 3import streamlit as st4from annotated_text import annotated_text5from cherche import compose, qa, rank, retrieve, summary6from sentence_transformers import SentenceTransformer7from sklearn.feature_extraction.text import TfidfVectorizer8from transformers import pipeline9 10 11@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)12def loading_pipelines():13 """Create three pipelines dedicated to neural research. The first one is dedicated to game14 retrieval. The second is dedicated to the question answering task. The third is dedicated to15 the summarization task. Save pipelines as pickle file.16 17 >>> search = (18 ... tfidf(on = "game") + ranker(on = "game") | tfidf(on = ["game", "summary"]) +19 ... ranker(on = ["game", "summary"]) + documents20 ... )21 22 """23 # Load documents24 with open("games.json", "r") as documents_file:25 documents = json.load(documents_file)26 27 # A first retriever dedicated to title28 retriever_title = retrieve.TfIdf(29 key="id",30 on=["game"],31 documents=documents,32 tfidf=TfidfVectorizer(33 lowercase=True,34 min_df=1,35 max_df=0.9,36 ngram_range=(3, 7),37 analyzer="char",38 ),39 k=30,40 )41 42 # A second retriever dedicated to title and also summary of games.43 retriever_title_summary = retrieve.TfIdf(44 key="id",45 on=["game", "summary"],46 documents=documents,47 tfidf=TfidfVectorizer(48 lowercase=True,49 min_df=1,50 max_df=0.9,51 ngram_range=(3, 7),52 analyzer="char",53 ),54 k=30,55 )56 57 # Load our encoder to re-rank retrievers documents.58 encoder = SentenceTransformer("sentence-transformers/all-mpnet-base-v2").encode59 60 # A ranker dedicated to title61 ranker_title = rank.Encoder(62 key="id",63 on=["game"],64 encoder=encoder,65 k=5,66 path="games_title.pkl",67 )68 69 # A ranker dedicated to title and summary70 ranker_title_summary = rank.Encoder(71 key="id",72 on=["game", "summary"],73 encoder=encoder,74 k=5,75 path="games_summary.pkl",76 )77 78 # Pipeline creation79 search = (80 (retriever_title + ranker_title) | (retriever_title_summary + ranker_title_summary)81 ) + documents82 83 # Index84 search.add(documents)85 return search86 87 88@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)89def write_search(query):90 return search(query)[:5]91 92 93@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)94def loading_summarization_pipeline():95 summarizer = summary.Summary(96 model=pipeline(97 "summarization",98 model="sshleifer/distilbart-cnn-12-6",99 tokenizer="sshleifer/distilbart-cnn-12-6",100 framework="pt",101 ),102 on=["game", "summary"],103 max_length=50,104 )105 106 search_summarize = search + summarizer107 return search_summarize108 109 110@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)111def write_search_summarize(query_summarize):112 return search_summarize(query_summarize)113 114 115@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)116def loading_qa_pipeline():117 question_answering = qa.QA(118 model=pipeline(119 "question-answering",120 model="deepset/roberta-base-squad2",121 tokenizer="deepset/roberta-base-squad2",122 ),123 k=3,124 on="summary",125 )126 search_qa = search + question_answering127 return search_qa128 129 130@st.cache(hash_funcs={compose.Pipeline: lambda _: None}, allow_output_mutation=True)131def write_search_qa(query_qa):132 return search_qa(query_qa)133 134 135if __name__ == "__main__":136 137 st.markdown("# 🕹 Cherche")138 139 st.markdown(140 "[Cherche](https://github.com/raphaelsty/cherche) (search in French) allows you to create a \141 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."142 )143 144 st.image("explain.png")145 146 st.markdown(147 "Here is a demo of neural search for video games using a sample of reviews made by [Metacritic](https://www.metacritic.com). \148 Starting the app may take a while if the models are not stored in cache."149 )150 151 # Will be slow the first time, you will need to compute embeddings.152 search = loading_pipelines()153 154 st.markdown("## 👾 Neural search")155 156 st.markdown(157 '```search = (tfidf(on = "title") + ranker(on = "title") | tfidf(on = ["title", "summary"]) + ranker(on = ["game", "summary"]) + documents)```'158 )159 160 query = st.text_input(161 "games",162 value="super smash bros",163 max_chars=None,164 key=None,165 type="default",166 help=None,167 autocomplete=None,168 on_change=None,169 args=None,170 kwargs=None,171 )172 173 if query:174 175 for document in write_search(query):176 if document["rate"] < 10:177 document["rate"] *= 10178 179 st.markdown(f"### {document['game']}")180 st.markdown(f"Metacritic Rating: {document['rate']}")181 182 col_1, col_2 = st.columns([1, 5])183 with col_1:184 st.image(document["image"], width=100)185 with col_2:186 st.write(f"{document['summary'][:430]}...")187 188 st.markdown("## 🎲 Summarization")189 190 st.markdown(191 '```search = (tfidf(on = "title") + ranker(on = "title") | tfidf(on = ["title", "summary"]) + ranker(on = ["game", "summary"]) + documents + summarization(on = "summary"))```'192 )193 194 st.markdown(195 "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.."196 )197 198 query_summarize = st.text_input(199 "summarization",200 value="super smash bros",201 max_chars=None,202 key=None,203 type="default",204 help=None,205 autocomplete=None,206 on_change=None,207 args=None,208 kwargs=None,209 )210 211 if query_summarize:212 search_summarize = loading_summarization_pipeline()213 st.write(f"**{write_search_summarize(query_summarize)}**")214 215 st.markdown("## 🎮 Question answering")216 217 st.markdown(218 '```search = (tfidf(on = "title") + ranker(on = "title") | tfidf(on = ["title", "summary"]) + ranker(on = ["game", "summary"]) + documents + question_answering(on = "summary"))```'219 )220 221 st.markdown(222 "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."223 )224 225 query_qa = st.text_input(226 "question",227 value="What is the purpose of playing Super Smash Bros?",228 max_chars=None,229 key=None,230 type="default",231 help=None,232 autocomplete=None,233 on_change=None,234 args=None,235 kwargs=None,236 )237 238 if query_qa:239 240 search_qa = loading_qa_pipeline()241 for document_qa in write_search_qa(query_qa):242 243 st.markdown(f"### {document_qa['game']}")244 st.markdown(f"Metacritic Rating: {document_qa['rate']}")245 246 col_1, col_2 = st.columns([1, 5])247 with col_1:248 st.image(document_qa["image"], width=100)249 with col_2:250 251 annotations = document_qa["summary"].split(document_qa["answer"])252 253 if document_qa["start"] == 0:254 annotated_text(255 (256 document_qa["answer"],257 f"answer {round(document_qa['qa_score'], 2)}",258 "#8ef",259 ),260 " ",261 " ".join(annotations[1:]),262 )263 264 elif document_qa["end"] == len(document_qa["summary"]):265 annotated_text(266 " ".join(annotations[:-1]),267 (268 document_qa["answer"],269 f"answer {round(document_qa['qa_score'], 2)}",270 "#8ef",271 ),272 )273 274 else:275 annotated_text(276 annotations[0],277 (278 document_qa["answer"],279 f"answer {round(document_qa['qa_score'], 2)}",280 "#8ef",281 ),282 annotations[1],283 )284