From 316b7d471a0344cbd47e779aa49f6ef431ed0f4f Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Mon, 4 Dec 2023 13:46:25 -0500 Subject: [PATCH] Handle offline chat model retrieval when no internet Offline chat shouldn't fail on retrieve_model when no internet, if model was previously downloaded and usable offline --- src/khoj/processor/conversation/offline/utils.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/khoj/processor/conversation/offline/utils.py b/src/khoj/processor/conversation/offline/utils.py index 0b876b26..3a1862f7 100644 --- a/src/khoj/processor/conversation/offline/utils.py +++ b/src/khoj/processor/conversation/offline/utils.py @@ -12,11 +12,12 @@ def download_model(model_name: str): logger.info("There was an error importing GPT4All. Please run pip install gpt4all in order to install it.") raise e - # Download the chat model - chat_model_config = gpt4all.GPT4All.retrieve_model(model_name=model_name, allow_download=True) - # Decide whether to load model to GPU or CPU + chat_model_config = None try: + # Download the chat model and its config + chat_model_config = gpt4all.GPT4All.retrieve_model(model_name=model_name, allow_download=True) + # Try load chat model to GPU if: # 1. Loading chat model to GPU isn't disabled via CLI and # 2. Machine has GPU @@ -26,6 +27,12 @@ def download_model(model_name: str): ) except ValueError: device = "cpu" + except Exception as e: + if chat_model_config is None: + device = "cpu" # Fallback to CPU as can't determine if GPU has enough memory + logger.debug(f"Unable to download model config from gpt4all website: {e}") + else: + raise e # Now load the downloaded chat model onto appropriate device chat_model = gpt4all.GPT4All(model_name=model_name, device=device, allow_download=False)