added command line chat

Makefile

@@ -10,7 +10,7 @@ else
 endif
 
 test:
-	PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 python test.py $(TEST)
+	python test.py $(TEST)
 
 chat:
 	python test.py chat

test.py

@@ -1,6 +1,7 @@
 # project/test.py
 
 import os
+import sys
 import unittest
 from timeit import default_timer as timer
 
@@ -134,5 +135,40 @@ class TestQAChain(unittest.TestCase):
         self.run_test_case("huggingface", self.question)
 
 
+def chat():
+    start = timer()
+    llm_loader = app_init()[0]
+    end = timer()
+    print(f"Model loaded in {end - start:.3f}s")
+
+    chat_chain = ChatChain(llm_loader)
+    chat_history = []
+
+    chat_start = timer()
+
+    while True:
+        query = input("Please enter your question: ")
+        query = query.strip()
+        if query.lower() == "exit":
+            break
+
+        print("\nQuestion: " + query)
+
+        start = timer()
+        result = chat_chain.call_chain(
+            {"question": query, "chat_history": chat_history}, None
+        )
+        end = timer()
+        print(f"Completed in {end - start:.3f}s")
+
+        chat_history.append((query, result["text"]))
+
+    chat_end = timer()
+    print(f"Total time used: {chat_end - chat_start:.3f}s")
+
+
 if __name__ == "__main__":
-    unittest.main()
+    if len(sys.argv) > 1 and sys.argv[1] == "chat":
+        chat()
+    else:
+        unittest.main()

test.sh

@@ -1,79 +0,0 @@
-#!/bin/sh
-
-EXT="$1"
-
-if [ "$EXT" = "" ]; then
-    echo usage: $0 log_ext
-    exit
-fi
-
-echo Using extension: $EXT
-
-[ ! -f .env ] || export $(grep -v '^#' .env | xargs)
-
-export LLM_MODEL_TYPE=openai
-export OPENAI_MODEL_NAME="gpt-3.5-turbo"
-echo Testing openai-${OPENAI_MODEL_NAME}
-python test.py 2>&1 | tee ./data/logs/openai-${OPENAI_MODEL_NAME}_${EXT}.log
-
-export OPENAI_MODEL_NAME="gpt-4"
-echo Testing openai-${OPENAI_MODEL_NAME}
-python test.py 2>&1 | tee ./data/logs/openai-${OPENAI_MODEL_NAME}_${EXT}.log
-
-export LLM_MODEL_TYPE=huggingface
-
-export HUGGINGFACE_MODEL_NAME_OR_PATH="lmsys/fastchat-t5-3b-v1.0"
-echo Testing $HUGGINGFACE_MODEL_NAME_OR_PATH
-python test.py 2>&1 | tee ./data/logs/fastchat-t5-3b-v1.0_${EXT}.log
-
-
-export HUGGINGFACE_MODEL_NAME_OR_PATH="TheBloke/wizardLM-7B-HF"
-echo Testing $HUGGINGFACE_MODEL_NAME_OR_PATH
-python test.py 2>&1 | tee ./data/logs/wizardLM-7B-HF_${EXT}.log
-
-
-export HUGGINGFACE_MODEL_NAME_OR_PATH="TheBloke/vicuna-7B-1.1-HF"
-echo Testing $HUGGINGFACE_MODEL_NAME_OR_PATH
-python test.py 2>&1 | tee ./data/logs/vicuna-7B-1.1-HF_${EXT}.log
-
-
-export HUGGINGFACE_MODEL_NAME_OR_PATH="nomic-ai/gpt4all-j"
-echo Testing $HUGGINGFACE_MODEL_NAME_OR_PATH
-python test.py 2>&1 | tee ./data/logs/gpt4all-j_${EXT}.log
-
-
-# export HUGGINGFACE_MODEL_NAME_OR_PATH="nomic-ai/gpt4all-falcon"
-# echo Testing $HUGGINGFACE_MODEL_NAME_OR_PATH
-# python test.py 2>&1 | tee ./data/logs/gpt4all-falcon_${EXT}.log
-
-export LLM_MODEL_TYPE=stablelm
-
-# export STABLELM_MODEL_NAME_OR_PATH="stabilityai/stablelm-tuned-alpha-7b"
-# echo Testing $STABLELM_MODEL_NAME_OR_PATH
-# python test.py 2>&1 | tee ./data/logs/stablelm-tuned-alpha-7b_${EXT}.log
-
-
-export STABLELM_MODEL_NAME_OR_PATH="OpenAssistant/stablelm-7b-sft-v7-epoch-3"
-echo Testing $STABLELM_MODEL_NAME_OR_PATH
-python test.py 2>&1 | tee ./data/logs/stablelm-7b-sft-v7-epoch-3_${EXT}.log
-
-
-export LLM_MODEL_TYPE=mosaicml
-export MOSAICML_MODEL_NAME_OR_PATH="mosaicml/mpt-7b-instruct"
-echo Testing $MOSAICML_MODEL_NAME_OR_PATH
-python test.py 2>&1 | tee ./data/logs/mpt-7b-instruct_${EXT}.log
-
-
-# export MOSAICML_MODEL_NAME_OR_PATH="mosaicml/mpt-30b-instruct"
-# echo Testing $MOSAICML_MODEL_NAME_OR_PATH
-# LOAD_QUANTIZED_MODEL=4bit python test.py 2>&1 | tee ./data/logs/mpt-30b-instruct_${EXT}.log
-
-export LLM_MODEL_TYPE=huggingface
-export HUGGINGFACE_MODEL_NAME_OR_PATH="HuggingFaceH4/starchat-beta"
-echo Testing $HUGGINGFACE_MODEL_NAME_OR_PATH
-LOAD_QUANTIZED_MODEL=8bit python test.py 2>&1 | tee ./data/logs/starchat-beta_${EXT}.log
-
-
-# export HUGGINGFACE_MODEL_NAME_OR_PATH="../../models/starcoder"
-# echo Testing $HUGGINGFACE_MODEL_NAME_OR_PATH
-# LOAD_QUANTIZED_MODEL=8bit python test.py 2>&1 | tee ./data/logs/starcoder_${EXT}.log