add option to select model

Author: mmz-001

knowledge_gpt/core/qa.py

@@ -1,11 +1,10 @@
-from typing import Any, List
+from typing import List
 from langchain.chains.qa_with_sources import load_qa_with_sources_chain
 from knowledge_gpt.core.prompts import STUFF_PROMPT
 from langchain.docstore.document import Document
-from langchain.chat_models import ChatOpenAI
 from knowledge_gpt.core.embedding import FolderIndex
-from knowledge_gpt.core.debug import FakeChatModel
 from pydantic import BaseModel
+from langchain.chat_models.base import BaseChatModel
 
 
 class AnswerWithSources(BaseModel):
@@ -16,9 +15,8 @@ class AnswerWithSources(BaseModel):
 def query_folder(
     query: str,
     folder_index: FolderIndex,
+    llm: BaseChatModel,
     return_all: bool = False,
-    model: str = "openai",
-    **model_kwargs: Any,
 ) -> AnswerWithSources:
     """Queries a folder index for an answer.
 
@@ -33,15 +31,6 @@ def query_folder(
     Returns:
         AnswerWithSources: The answer and the source documents.
     """
-    supported_models = {
-        "openai": ChatOpenAI,
-        "debug": FakeChatModel,
-    }
-
-    if model in supported_models:
-        llm = supported_models[model](**model_kwargs)
-    else:
-        raise ValueError(f"Model {model} not supported.")
 
     chain = load_qa_with_sources_chain(
         llm=llm,
@@ -73,5 +62,4 @@ def get_sources(answer: str, folder_index: FolderIndex) -> List[Document]:
         for doc in file.docs:
             if doc.metadata["source"] in source_keys:
                 source_docs.append(doc)
-
     return source_docs

knowledge_gpt/core/utils.py

@@ -2,6 +2,10 @@
 from langchain.chains.combine_documents.stuff import StuffDocumentsChain
 from langchain.docstore.document import Document
 
+from langchain.chat_models import ChatOpenAI
+from knowledge_gpt.core.debug import FakeChatModel
+from langchain.chat_models.base import BaseChatModel
+
 
 def pop_docs_upto_limit(
     query: str, chain: StuffDocumentsChain, docs: List[Document], max_len: int
@@ -16,3 +20,13 @@ def pop_docs_upto_limit(
         token_count = chain.prompt_length(docs, question=query)  # type: ignore
 
     return docs
+
+
+def get_llm(model: str, **kwargs) -> BaseChatModel:
+    if model == "debug":
+        return FakeChatModel()
+
+    if "gpt" in model:
+        return ChatOpenAI(model=model, **kwargs)  # type: ignore
+
+    raise ValueError(f"Model {model} not supported!")

knowledge_gpt/main.py

@@ -16,13 +16,15 @@
 from knowledge_gpt.core.chunking import chunk_file
 from knowledge_gpt.core.embedding import embed_files
 from knowledge_gpt.core.qa import query_folder
+from knowledge_gpt.core.utils import get_llm
+
 
 EMBEDDING = "openai"
 VECTOR_STORE = "faiss"
-MODEL = "openai"
+MODEL_LIST = ["gpt-3.5-turbo", "gpt-4"]
 
-# For testing
-EMBEDDING, VECTOR_STORE, MODEL = ["debug"] * 3
+# Uncomment to enable debug mode
+# MODEL_LIST.insert(0, "debug")
 
 st.set_page_config(page_title="KnowledgeGPT", page_icon="📖", layout="wide")
 st.header("📖KnowledgeGPT")
@@ -48,6 +50,13 @@
     help="Scanned documents are not supported yet!",
 )
 
+model: str = st.selectbox("Model", options=MODEL_LIST)  # type: ignore
+
+with st.expander("Advanced Options"):
+    return_all_chunks = st.checkbox("Show all chunks retrieved from vector search")
+    show_full_doc = st.checkbox("Show parsed contents of the document")
+
+
 if not uploaded_file:
     st.stop()
 
@@ -61,15 +70,16 @@
 if not is_file_valid(file):
     st.stop()
 
-if MODEL != "debug" and not is_open_ai_key_valid(openai_api_key):
+
+if not is_open_ai_key_valid(openai_api_key, model):
     st.stop()
 
 
 with st.spinner("Indexing document... This may take a while⏳"):
     folder_index = embed_files(
         files=[chunked_file],
-        embedding=EMBEDDING,
-        vector_store=VECTOR_STORE,
+        embedding=EMBEDDING if model != "debug" else "debug",
+        vector_store=VECTOR_STORE if model != "debug" else "debug",
         openai_api_key=openai_api_key,
     )
 
@@ -78,11 +88,6 @@
     submit = st.form_submit_button("Submit")
 
 
-with st.expander("Advanced Options"):
-    return_all_chunks = st.checkbox("Show all chunks retrieved from vector search")
-    show_full_doc = st.checkbox("Show parsed contents of the document")
-
-
 if show_full_doc:
     with st.expander("Document"):
         # Hack to get around st.markdown rendering LaTeX
@@ -96,13 +101,12 @@
     # Output Columns
     answer_col, sources_col = st.columns(2)
 
+    llm = get_llm(model=model, openai_api_key=openai_api_key, temperature=0)
     result = query_folder(
         folder_index=folder_index,
         query=query,
         return_all=return_all_chunks,
-        model=MODEL,
-        openai_api_key=openai_api_key,
-        temperature=0,
+        llm=llm,
     )
 
     with answer_col:

knowledge_gpt/ui.py

@@ -43,18 +43,22 @@ def display_file_read_error(e: Exception) -> NoReturn:
 
 
 @st.cache_data(show_spinner=False)
-def is_open_ai_key_valid(openai_api_key) -> bool:
+def is_open_ai_key_valid(openai_api_key, model: str) -> bool:
+    if model == "debug":
+        return True
+
     if not openai_api_key:
         st.error("Please enter your OpenAI API key in the sidebar!")
         return False
     try:
         openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
+            model=model,
             messages=[{"role": "user", "content": "test"}],
             api_key=openai_api_key,
         )
     except Exception as e:
         st.error(f"{e.__class__.__name__}: {e}")
         logger.error(f"{e.__class__.__name__}: {e}")
         return False
+
     return True