revert: remove option to use `togethercomputer/GPT-NeoXT-Chat-Base-20B`

app.py

@@ -1,6 +1,4 @@
 import os
-from typing import Literal
-import logging
 import streamlit as st
 
 from langchain.embeddings import HuggingFaceInstructEmbeddings
@@ -9,7 +7,6 @@ from langchain.chains import VectorDBQA
 from huggingface_hub import snapshot_download
 from langchain import OpenAI
 from langchain import PromptTemplate
-from langchain.llms import HuggingFacePipeline, HuggingFaceHub
 
 
 BOOK_NAME = "1984"
@@ -76,26 +73,10 @@ def load_prompt(book_name, author_name):
     return PROMPT
 
     
-@st.experimental_singleton(show_spinner=False, max_entries=1)
-def load_chain(model: Literal["openai", "togethercomputer/GPT-NeoXT-Chat-Base-20B"] ="openai"):
-
-    # choose model
-    if model=="openai":
-        llm = OpenAI(temperature=0.2)
-
-    if model=="togethercomputer/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}
-        # )
-        llm = HuggingFaceHub(
-            repo_id="togethercomputer/GPT-NeoXT-Chat-Base-20B", 
-            task="text-generation",
-            model_kwargs={"temperature":0.2, "max_length":400}
-        )        
-
-    # load chain
+@st.experimental_singleton(show_spinner=False)
+def load_chain():
+    llm = OpenAI(temperature=0.2)
+
     chain = VectorDBQA.from_chain_type(
         chain_type_kwargs = {"prompt": load_prompt(book_name=BOOK_NAME, author_name=AUTHOR_NAME)},
         llm=llm,
@@ -104,14 +85,12 @@ def load_chain(model: Literal["openai", "togethercomputer/GPT-NeoXT-Chat-Base-20
         k=8,
         return_source_documents=True,
         )
-
-    logging.info(f"Loaded chain with {model}.")
     
     return chain
 
 
-def get_answer(question, model="openai"):
-    chain = load_chain(model=model)
+def get_answer(question):
+    chain = load_chain()
     result = chain({"query": question})
 
     answer = result["result"]
@@ -145,26 +124,11 @@ def get_answer(question, model="openai"):
 
 ##### sidebar ####
 with st.sidebar:
-    
-    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
-
-    
-    if choice == "togethercomputer/GPT-NeoXT-Chat-Base-20B":
-        api_key = st.text_input(label = "Paste your Hugging Face Hub API key here to get started", 
-                                type = "password",
-                                help = "This isn't saved 🙈"
-                               )
-        os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
+    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("---")
 
@@ -187,21 +151,16 @@ ask = col2.button("Ask")
 
 if ask:
 
-    if (choice=="openai" and os.environ["OPENAI_API_KEY"]=="") or (choice=="togethercomputer/GPT-NeoXT-Chat-Base-20B" and os.environ["HUGGINGFACEHUB_API_TOKEN"]==""):
+    if 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, model=choice)
-                logging.info(f"Answer successfully generated using {choice}.")
+                answer, pages, extract = get_answer(question=user_input)
             except:
-                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}:** What\'s going on? That's not the right API key")
+                st.stop()
 
     st.write(f"**{BOOK_NAME}:** {answer}")