add option to use `togethercomputer/GPT-NeoXT-Chat-Base-20B`

app.py

@@ -1,4 +1,5 @@
 import os
+from typing import Literal
 import streamlit as st
 
 from langchain.embeddings import HuggingFaceInstructEmbeddings
@@ -7,6 +8,7 @@ from langchain.chains import VectorDBQA
 from huggingface_hub import snapshot_download
 from langchain import OpenAI
 from langchain import PromptTemplate
+from langchain.llms import HuggingFacePipeline
 
 
 BOOK_NAME = "1984"
@@ -74,9 +76,20 @@ def load_prompt(book_name, author_name):
 
     
 @st.experimental_singleton(show_spinner=False)
-def load_chain():
-    llm = OpenAI(temperature=0.2)
+def load_chain(model: Literal["openai", "GPT-NeoXT-Chat-Base-20B"] ="openai"):
 
+    # choose model
+    if model=="openai":
+        llm = OpenAI(temperature=0.2)
+
+    if model=="GPT-NeoXT-Chat-Base-20B":
+        llm = HuggingFacePipeline.from_model_id(
+            model_id="togethercomputer/GPT-NeoXT-Chat-Base-20B", 
+            task="text-generation",
+            model_kwargs={"temperature":0.2, "max_length":400}
+        )
+
+    # load chain
     chain = VectorDBQA.from_chain_type(
         chain_type_kwargs = {"prompt": load_prompt(book_name=BOOK_NAME, author_name=AUTHOR_NAME)},
         llm=llm,
@@ -89,8 +102,8 @@ def load_chain():
     return chain
 
 
-def get_answer(question):
-    chain = load_chain()
+def get_answer(question, model="openai"):
+    chain = load_chain(model=model)
     result = chain({"query": question})
 
     answer = result["result"]
@@ -124,11 +137,18 @@ def get_answer(question):
 
 ##### sidebar ####
 with st.sidebar:
-    api_key = st.text_input(label = "Paste your OpenAI API key here to get started", 
-                            type = "password",
-                            help = "This isn't saved 🙈"
-                           )
-    os.environ["OPENAI_API_KEY"] = api_key
+    
+    choice= st.radio("Choose your API:", 
+                     ["OpenAI", "togethercomputer/GPT-NeoXT-Chat-Base-20B"], 
+                     help="GPT-NeoXT-Chat-Base-20B doesn't need an API Key"
+                    )
+
+    if choice == "OpenAI":
+        api_key = st.text_input(label = "Paste your OpenAI API key here to get started", 
+                                type = "password",
+                                help = "This isn't saved 🙈"
+                               )
+        os.environ["OPENAI_API_KEY"] = api_key
 
     st.markdown("---")
 
@@ -151,16 +171,20 @@ ask = col2.button("Ask")
 
 if ask:
 
-    if api_key is "":
+    if choice=="OpenAI" and api_key is "":
         st.write(f"**{BOOK_NAME}:** Whoops looks like you forgot your API key buddy")
         st.stop()
     else:
         with st.spinner("Um... excuse me but... this can take about a minute for your first question because some stuff have to be downloaded 🥺👉🏻👈🏻"):
             try:
-                answer, pages, extract = get_answer(question=user_input)
+                answer, pages, extract = get_answer(question=user_input, model=choice)
             except:
-                st.write(f"**{BOOK_NAME}:** What\'s going on? That's not the right API key")
-                st.stop()
+                if choice=="togethercomputer/GPT-NeoXT-Chat-Base-20B":
+                    st.write("The model probably timed out :(")
+                    st.stop()
+                else:
+                    st.write(f"**{BOOK_NAME}:** What\'s going on? That's not the right API key")
+                    st.stop()
 
     st.write(f"**{BOOK_NAME}:** {answer}")