import streamlit as st from sentence_transformers import SentenceTransformer, util from transformers import (AutoModelForQuestionAnswering, AutoTokenizer, pipeline) import pandas as pd import regex as re from urllib import request # Select model for question answering model_name = "deepset/roberta-base-squad2" # Load model & tokenizer model = AutoModelForQuestionAnswering.from_pretrained(model_name) tokenizer = AutoTokenizer.from_pretrained(model_name) # Create pipeline pipe = pipeline('question-answering', model=model_name, tokenizer=model_name) # Load Harry Potter book corpus from link url = ("https://raw.githubusercontent.com/formcept/whiteboard/master/nbviewer/notebooks/data/harrypotter/Book%201%20-%20The%20Philosopher's%20Stone.txt") response = request.urlopen(url) book1_raw_0 = response.read().decode('utf8') # Text pre-processing # Remove page statements book1_raw_1 = re.sub(r'Page \| [0-9]+ Harry Potter [a-zA-Z \-]+J.K. Rowling', '', book1_raw_0) # Remove newlines book1_raw_1 = re.sub(r'\n', '', book1_raw_1) # Remove periods; this will relevant in the regrouping later book1_raw_1 = re.sub(r'Mr. ', 'Mr ', book1_raw_1) book1_raw_1 = re.sub(r'Ms. ', 'Ms ', book1_raw_1) book1_raw_1 = re.sub(r'Mrs. ', 'Mrs ', book1_raw_1) # Group into 3 sentences-long parts paragraphs = re.findall("[^.?!]+[.?!][^.?!]+[.?!][^.?!]+[.?!]", book1_raw_1) # Type in HP-related query here query = st.text_area("Hello muggle! What is your question?") # Perform sentence embedding on query and sentence groups model_embed_name = 'sentence-transformers/multi-qa-MiniLM-L6-cos-v1' model_embed = SentenceTransformer(model_embed_name) doc_emb = model_embed.encode(paragraphs) query_emb = model_embed.encode(query) #Compute dot score between query and all document embeddings scores = util.cos_sim(query_emb, doc_emb)[0].cpu().tolist() #Combine docs & scores doc_score_pairs = list(zip(paragraphs, scores)) #Sort by decreasing score and get only 3 most similar groups doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)[:3] # Join these similar groups to form the context context = "".join(x[0] for x in doc_score_pairs) # Perform the querying QA_input = {'question': query, 'context': context} out = pipe(QA_input) st.json(out)