Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| """wiki_chat.ipynb | |
| Automatically generated by Colaboratory. | |
| Original file is located at | |
| https://colab.research.google.com/drive/1P5rJeCXRSsDJw_1ksnHmodH6ng2Ot5NW | |
| """ | |
| # !pip install gradio | |
| # !pip install -U sentence-transformers | |
| # !pip install datasets | |
| import gradio as gr | |
| from sentence_transformers import SentenceTransformer, CrossEncoder, util | |
| from torch import tensor as torch_tensor | |
| from datasets import load_dataset | |
| """# import models""" | |
| bi_encoder = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1') | |
| bi_encoder.max_seq_length = 256 #Truncate long passages to 256 tokens | |
| #The bi-encoder will retrieve top_k documents. We use a cross-encoder, to re-rank the results list to improve the quality | |
| cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') | |
| """# import datasets""" | |
| dataset = load_dataset("gfhayworth/wiki_mini", split='train') | |
| mypassages = list(dataset.to_pandas()['psg']) | |
| dataset_embed = load_dataset("gfhayworth/wiki_mini_embed", split='train') | |
| dataset_embed_pd = dataset_embed.to_pandas() | |
| mycorpus_embeddings = torch_tensor(dataset_embed_pd.values) | |
| def search(query, top_k=20, top_n = 1): | |
| question_embedding = bi_encoder.encode(query, convert_to_tensor=True) | |
| question_embedding = question_embedding #.cuda() | |
| hits = util.semantic_search(question_embedding, mycorpus_embeddings, top_k=top_k) | |
| hits = hits[0] # Get the hits for the first query | |
| ##### Re-Ranking ##### | |
| cross_inp = [[query, mypassages[hit['corpus_id']]] for hit in hits] | |
| cross_scores = cross_encoder.predict(cross_inp) | |
| # Sort results by the cross-encoder scores | |
| for idx in range(len(cross_scores)): | |
| hits[idx]['cross-score'] = cross_scores[idx] | |
| hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True) | |
| predictions = hits[:top_n] | |
| return predictions | |
| # for hit in hits[0:3]: | |
| # print("\t{:.3f}\t{}".format(hit['cross-score'], mypassages[hit['corpus_id']].replace("\n", " "))) | |
| def get_text(qry): | |
| predictions = search(qry) | |
| prediction_text = [] | |
| for hit in predictions: | |
| prediction_text.append("{}".format(mypassages[hit['corpus_id']])) | |
| return prediction_text | |
| # def prt_rslt(qry): | |
| # rslt = get_text(qry) | |
| # for r in rslt: | |
| # print(r) | |
| # prt_rslt("who is the best rapper in the world?") | |
| """# chat example""" | |
| def chat(message, history): | |
| history = history or [] | |
| message = message.lower() | |
| responses = get_text(message) | |
| for response in responses: | |
| history.append((message, response)) | |
| return history, history | |
| css=".gradio-container {background-color: lightgray}" | |
| with gr.Blocks(css=css) as demo: | |
| history_state = gr.State() | |
| gr.Markdown('# WikiBot') | |
| title='Wikipedia Chatbot' | |
| description='chatbot with search on Wikipedia' | |
| with gr.Row(): | |
| chatbot = gr.Chatbot() | |
| with gr.Row(): | |
| message = gr.Textbox(label='Input your question here:', | |
| placeholder='How many countries are in Europe?', | |
| lines=1) | |
| submit = gr.Button(value='Send', | |
| variant='secondary').style(full_width=False) | |
| submit.click(chat, | |
| inputs=[message, history_state], | |
| outputs=[chatbot, history_state]) | |
| gr.Examples( | |
| examples=["How many countries are in Europe?", | |
| "Was Roman Emperor Constantine I a Christian?", | |
| "Who is the best rapper in the world?"], | |
| inputs=message | |
| ) | |
| demo.launch() | |