Add main application

app.py

@@ -7,14 +7,77 @@ import nltk
 from nltk import sent_tokenize
 nltk.download("punkt")
 
-krishnamurti_generator =  pipeline("text-generation", model="eliwill/stoic-generator-10e")
-# stoic_generator = pipeline("text-generation", model="eliwill/stoic-generator-10e")
+# Loading in dataframes
+krishnamurti_df = pd.read_json("krishnamurti_df.json")
+stoic_df = pd.read_json("stoic_df.json")
 
- ############### BUILDING DEMO ################################
+# 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['quote'][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)
+  
 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()