Update app.py

app.py

@@ -2,6 +2,7 @@ import numpy as np
 import requests
 import streamlit as st
 import openai
+import json
 
 def main():
     st.title("Scientific Question Generation")
@@ -12,16 +13,27 @@ def main():
     of the [FLAN-T5-XXL](https://huggingface.co/google/flan-t5-xxl) model and the [GPT-3.5-turbo](https://platform.openai.com/docs/models/gpt-3-5) model.\
     \n\n For a more thorough discussion of question generation see this [report](https://wandb.ai/dmeltzer/Question_Generation/reports/Exploratory-Data-Analysis-for-r-AskScience--Vmlldzo0MjQwODg1?accessToken=fndbu2ar26mlbzqdphvb819847qqth2bxyi4hqhugbnv97607mj01qc7ed35v6w8) for EDA on the r/AskScience dataset and this \
     [report](https://api.wandb.ai/links/dmeltzer/7an677es) for details on our training procedure.\
-    \n\n**Disclaimer**: You may recieve an error message when you first run the model. We are using the Huggingface API to access the BART-Large and FLAN-T5 models, and the inference API takes around 20 seconds to load each model.\
+    \n\n \
+    The two fine-tuned models (BART-Large and FLAN-T5-Base) are hosted on AWS using a combination of AWS Sagemaker, Lambda, and API gateway. \
+    \ GPT-3.5 is called using the OpenAI API and the FLAN-T5-XXL model is hosted by HuggingFace and is called with their Inference API.\
+    \n \n \
+    **Disclaimer**: You may recieve an error message when calling the FLAN-T5-XXL model since the Inference API takes around 20 seconds to load the model.\
     ")
     
-    checkpoints = ['dhmeltzer/bart-large_askscience-qg',
-              'dhmeltzer/flan-t5-base_askscience-qg',
-              'google/flan-t5-xxl']
+    AWS_checkpoints = {}
+    AWS_checkpoints['BART-Large']='https://8hlnvys7bh.execute-api.us-east-1.amazonaws.com/beta/'
+    AWS_checkpoints['FLAN-T5-Base']='https://gnrxh05827.execute-api.us-east-1.amazonaws.com/beta/'
     
-    headers = {"Authorization": f"Bearer {st.secrets['HF_token']}"}
+    # Right now HF_checkpoints just consists of FLAN-T5-XXL but we may add more models later.
+    HF_checkpoints = ['google/flan-t5-xxl']
+
+    # Token to access HF inference API
+    HF_headers = {"Authorization": f"Bearer {st.secrets['HF_token']}"}
+
+    # Token to access OpenAI API
     openai.api_key = st.secrets['OpenAI_token']
-    
+
+    # Used to query models hosted on Huggingface
     def query(checkpoint, payload):
         API_URL = f"https://api-inference.huggingface.co/models/{checkpoint}"
         
@@ -36,12 +48,41 @@ def main():
                                 """Black holes are the most gravitationally dense objects in the universe.""")
     
     if user_input:
-        for checkpoint in checkpoints:
+
+        for name, url in AWS_checkpoints.values():
+            headers={'x-api-key': key}
+
+            input_data = json.dumps({'inputs':user_input})
+            r = requests.get(url,data=input_data,headers=headers)
+            output = r.json()[0]['generated_text']
+
+            st.write(f'**{name}**: {output}')
+
+        model_engine = "gpt-3.5-turbo"
+        # Max tokens to produce
+        max_tokens = 50
+
+        # Prompt GPT-3.5 with an explicit question
+        prompt = f"generate a question: {user_input}"
+
+        # We give GPT-3.5 a message so it knows to generate questions from text.
+        response=openai.ChatCompletion.create(
+            model=model_engine,
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant that generates questions from text."},
+                {"role": "user", "content": prompt},
+            ])
+    
+        output = response['choices'][0]['message']['content']   
+        st.write(f'**{model_engine}**: {output}')
+
+        
+        for checkpoint in HF_checkpoints:
             
             model_name = checkpoint.split('/')[1]
     
+            # For FLAN models we need to give them instructions explicitly.
             if 'flan' in model_name.lower():
-                
                 prompt = 'generate a question: ' + user_input
 
             else:
@@ -57,23 +98,6 @@ def main():
                 return
             
             st.write(f'**{model_name}**: {output}')
-    
-        model_engine = "gpt-3.5-turbo"
-        max_tokens = 50
-        
-        prompt = f"generate a question: {user_input}"
-    
-        response=openai.ChatCompletion.create(
-            model=model_engine,
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant that generates questions from text."},
-                {"role": "user", "content": prompt},
-            ])
-    
-        output = response['choices'][0]['message']['content']
-        
-        st.write(f'**{model_engine}**: {output}')
-
         
 if __name__ == "__main__":
     main()