Add chatdocs.yml configuration file

Replace all command-line options with the new `chatdocs.yml` config

Author: marella

.github/workflows/tests.yml

@@ -35,8 +35,14 @@ jobs:
           python -m pip install --upgrade pip
           pip install .
 
+      - name: Copy chatdocs.yml
+        run: cp tests/fixtures/chatdocs.yml .
+
+      - name: Test download
+        run: chatdocs download
+
       - name: Test add
         run: chatdocs add examples
 
       - name: Test chat
-        run: chatdocs chat 'Why was the NATO created?' --model marella/gpt-2-ggml --download --lib avx
+        run: chatdocs chat 'Why was the NATO created?'

.gitignore

@@ -1,3 +1,4 @@
+/chatdocs.yml
 /db/
 
 # Created by https://www.toptal.com/developers/gitignore/api/c++,python,cmake,linux,macos,windows,sublimetext,vim,visualstudio,visualstudiocode

chatdocs/add.py

@@ -1,6 +1,6 @@
 import os
 import glob
-from typing import List
+from typing import Any, Dict, List
 from multiprocessing import Pool
 
 from tqdm import tqdm
@@ -18,12 +18,9 @@
     UnstructuredWordDocumentLoader,
 )
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.docstore.document import Document
-from chromadb.config import Settings
 
-from . import config
+from .vectorstores import get_vectorstore, get_vectorstore_from_documents
 
 
 # Custom document loaders
@@ -143,23 +140,12 @@ def does_vectorstore_exist(persist_directory: str) -> bool:
     return False
 
 
-def add(source_directory: str, persist_directory: str) -> None:
-    # Create embeddings
-    embeddings = HuggingFaceInstructEmbeddings(model_name=config.EMBEDDINGS_MODEL)
-    chroma_settings = Settings(
-        chroma_db_impl=config.CHROMA_DB_IMPL,
-        persist_directory=persist_directory,
-        anonymized_telemetry=False,
-    )
-
+def add(config: Dict[str, Any], source_directory: str) -> None:
+    persist_directory = config["chroma"]["persist_directory"]
     if does_vectorstore_exist(persist_directory):
         # Update and store locally vectorstore
         print(f"Appending to existing vectorstore at {persist_directory}")
-        db = Chroma(
-            persist_directory=persist_directory,
-            embedding_function=embeddings,
-            client_settings=chroma_settings,
-        )
+        db = get_vectorstore(config)
         collection = db.get()
         texts = process_documents(
             source_directory,
@@ -172,11 +158,6 @@ def add(source_directory: str, persist_directory: str) -> None:
         print("Creating new vectorstore")
         texts = process_documents(source_directory)
         print(f"Creating embeddings. May take a few minutes...")
-        db = Chroma.from_documents(
-            texts,
-            embeddings,
-            persist_directory=persist_directory,
-            client_settings=chroma_settings,
-        )
+        db = get_vectorstore_from_documents(config, texts)
     db.persist()
     db = None

chatdocs/chains.py

@@ -0,0 +1,17 @@
+from typing import Any, Dict
+
+from langchain.chains import RetrievalQA
+
+from .llms import get_llm
+from .vectorstores import get_vectorstore
+
+
+def get_retrieval_qa(config: Dict[str, Any]) -> RetrievalQA:
+    db = get_vectorstore(config)
+    retriever = db.as_retriever(**config["retriever"])
+    llm = get_llm(config)
+    return RetrievalQA.from_chain_type(
+        llm=llm,
+        retriever=retriever,
+        return_source_documents=True,
+    )

chatdocs/chat.py

@@ -1,75 +1,15 @@
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
-from chromadb.config import Settings
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-from langchain.chains import RetrievalQA
-from langchain.embeddings import HuggingFaceInstructEmbeddings
-from langchain.llms import CTransformers, HuggingFacePipeline
-from langchain.vectorstores import Chroma
 from rich import print
 from rich.markup import escape
 from rich.panel import Panel
 
-from . import config
+from .chains import get_retrieval_qa
+from .utils import print_answer
 
 
-def print_response(text: str) -> None:
-    print(f"[bright_cyan]{escape(text)}", end="", flush=True)
-
-
-class StreamingPrintCallbackHandler(StreamingStdOutCallbackHandler):
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        print_response(token)
-
-
-def chat(
-    *,
-    persist_directory: str,
-    hf: bool,
-    download: bool,
-    model: str,
-    model_type: Optional[str] = None,
-    model_file: Optional[str] = None,
-    lib: Optional[str] = None,
-    query: Optional[str] = None,
-) -> None:
-    local_files_only = not download
-    embeddings = HuggingFaceInstructEmbeddings(model_name=config.EMBEDDINGS_MODEL)
-    chroma_settings = Settings(
-        chroma_db_impl=config.CHROMA_DB_IMPL,
-        persist_directory=persist_directory,
-        anonymized_telemetry=False,
-    )
-    db = Chroma(
-        persist_directory=persist_directory,
-        embedding_function=embeddings,
-        client_settings=chroma_settings,
-    )
-    retriever = db.as_retriever(search_kwargs={"k": 4})
-
-    if hf:
-        llm = HuggingFacePipeline.from_model_id(
-            model_id=model,
-            task="text-generation",
-            model_kwargs={"local_files_only": local_files_only},
-            pipeline_kwargs={"max_new_tokens": 256},
-        )
-    else:
-        llm = CTransformers(
-            model=model,
-            model_type=model_type,
-            model_file=model_file,
-            config={"context_length": 1024, "local_files_only": local_files_only},
-            lib=lib,
-            callbacks=[StreamingPrintCallbackHandler()],
-        )
-
-    qa = RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=retriever,
-        return_source_documents=True,
-    )
+def chat(config: Dict[str, Any], query: Optional[str] = None) -> None:
+    qa = get_retrieval_qa(config)
 
     interactive = not query
     print()
@@ -89,8 +29,8 @@ def chat(
         print("[bold]A:", end="", flush=True)
 
         res = qa(query)
-        if hf:
-            print_response(res["result"])
+        if config["llm"] != "ctransformers":
+            print_answer(res["result"])
 
         print()
         for doc in res["source_documents"]:

chatdocs/config.py

@@ -1,6 +1,26 @@
-EMBEDDINGS_MODEL = "hkunlp/instructor-large"
-MODEL = "TheBloke/Wizard-Vicuna-7B-Uncensored-GGML"
-MODEL_TYPE = "llama"
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
 
-CHROMA_DB_IMPL = "duckdb+parquet"
-PERSIST_DIRECTORY = "db"
+import yaml
+
+from .utils import merge
+
+FILENAME = "chatdocs.yml"
+
+
+def _get_config(path: Union[Path, str]) -> Dict[str, Any]:
+    path = Path(path)
+    if path.is_dir():
+        path = path / FILENAME
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def get_config(path: Optional[Union[Path, str]] = None) -> Dict[str, Any]:
+    default_config = _get_config(Path(__file__).parent / "data")
+    if path is None:
+        path = Path() / FILENAME
+        if not path.is_file():
+            return default_config
+    config = _get_config(path)
+    return merge(default_config, config)

chatdocs/data/chatdocs.yml

@@ -0,0 +1,26 @@
+embeddings:
+  model: hkunlp/instructor-large
+
+llm: ctransformers
+
+ctransformers:
+  model: TheBloke/Wizard-Vicuna-7B-Uncensored-GGML
+  model_type: llama
+  config:
+    context_length: 1024
+
+huggingface:
+  model: TheBloke/Wizard-Vicuna-7B-Uncensored-HF
+  pipeline_kwargs:
+    max_new_tokens: 256
+
+download: false
+
+chroma:
+  persist_directory: db
+  chroma_db_impl: duckdb+parquet
+  anonymized_telemetry: false
+
+retriever:
+  search_kwargs:
+    k: 4

chatdocs/download.py

@@ -1,9 +1,10 @@
-from langchain.embeddings import HuggingFaceInstructEmbeddings
-from langchain.llms import CTransformers
+from typing import Any, Dict
 
-from . import config
+from .embeddings import get_embeddings
+from .llms import get_llm
 
 
-def download() -> None:
-    HuggingFaceInstructEmbeddings(model_name=config.EMBEDDINGS_MODEL)
-    CTransformers(model=config.MODEL, model_type=config.MODEL_TYPE)
+def download(config: Dict[str, Any]) -> None:
+    config = {**config, "download": True}
+    get_embeddings(config)
+    get_llm(config)

chatdocs/embeddings.py

@@ -0,0 +1,14 @@
+from typing import Any, Dict
+
+from langchain.embeddings import HuggingFaceInstructEmbeddings, HuggingFaceEmbeddings
+from langchain.embeddings.base import Embeddings
+
+
+def get_embeddings(config: Dict[str, Any]) -> Embeddings:
+    config = {**config["embeddings"]}
+    config["model_name"] = config.pop("model")
+    if config["model_name"].startswith("hkunlp/"):
+        Provider = HuggingFaceInstructEmbeddings
+    else:
+        Provider = HuggingFaceEmbeddings
+    return Provider(**config)

chatdocs/llms.py

@@ -0,0 +1,26 @@
+from typing import Any, Dict
+
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.llms import CTransformers, HuggingFacePipeline
+from langchain.llms.base import LLM
+
+from .utils import merge, print_answer
+
+
+class StreamingPrintCallbackHandler(StreamingStdOutCallbackHandler):
+    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        print_answer(token)
+
+
+def get_llm(config: Dict[str, Any]) -> LLM:
+    local_files_only = not config["download"]
+    if config["llm"] == "ctransformers":
+        config = {**config["ctransformers"]}
+        config = merge(config, {"config": {"local_files_only": local_files_only}})
+        llm = CTransformers(callbacks=[StreamingPrintCallbackHandler()], **config)
+    else:
+        config = {**config["huggingface"]}
+        config["model_id"] = config.pop("model")
+        config = merge(config, {"model_kwargs": {"local_files_only": local_files_only}})
+        llm = HuggingFacePipeline.from_model_id(task="text-generation", **config)
+    return llm