eliwill's picture
Add main application
de161f5
raw
history blame
3.06 kB
import gradio as gr
from transformers import pipeline
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer, util
import nltk
from nltk import sent_tokenize
nltk.download("punkt")
# Loading in dataframes
krishnamurti_df = pd.read_json("krishnamurti_df.json")
stoic_df = pd.read_json("stoic_df.json")
# Loading in sentence_similarity model
sentence_similarity_model = "all-mpnet-base-v2"
model = SentenceTransformer(sentence_similarity_model)
# Loading in text-generation models
stoic_generator = pipeline("text-generation", model="eliwill/stoic-generator-10e")
krishnamurti_generator = pipeline("text-generation", model="distilgpt2")
# Creating philosopher dictionary
philosopher_dictionary = {
"stoic": {
"generator": stoic_generator,
"dataframe": stoic_df
},
"krishnamurti": {
"generator": krishnamurti_generator,
"dataframe": krishnamurti_df
}
}
############### DEFINING FUNCTIONS ###########################
def ask_philosopher(philosopher, question):
""" Return first 5 sentences generated by question for the given philosopher model """
generator = philosopher_dictionary[philosopher]['generator']
answer = generator(question, min_length=100, max_length=120)[0]['generated_text'] # generate about 50 word tokens
answer = " ".join(sent_tokenize(answer)[:6]) # Get the first five sentences
return answer
def get_similar_quotes(philosopher, question):
""" Return top 5 most similar quotes to the question from a philosopher's dataframe """
df = philosopher_dictionary[philosopher]['dataframe']
question_embedding = model.encode(question)
sims = [util.dot_score(question_embedding, quote_embedding) for quote_embedding in df['Embedding']]
ind = np.argpartition(sims, -5)[-5:]
similar_sentences = [df['Quotes'][i] for i in ind]
top5quotes = pd.DataFrame(data = similar_sentences, columns=["Quotes"], index=range(1,6))
top5quotes['Quotes'] = top5quotes['Quotes'].str[:-1].str[:250] + "..."
return top5quotes
def main(question, philosopher):
return ask_philosopher(philosopher, question), get_similar_quotes(philosopher, question)
############### BUILDING DEMO ################################
with gr.Blocks() as demo:
gr.Markdown("""
# Ask a Philsopher
"""
)
with gr.Row():
with gr.Column():
inp1 = gr.Textbox(placeholder="Place your question here...", label="Ask a question")
inp2 = gr.Dropdown(choices=["stoic", "krishnamurti"], value="stoic", label="Choose a philosopher")
with gr.Column():
out1 = gr.Textbox(
lines=3,
max_lines=10,
label="Answer"
)
out2 = gr.DataFrame(
headers=["Quotes"],
max_rows=5,
interactive=False,
wrap=True)
btn = gr.Button("Run")
btn.click(fn=main, inputs=[inp1,inp2], outputs=[out1,out2])
demo.launch()