diff --git a/.gitignore b/.gitignore index 0fea4f7..fa326cf 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ .venv/ +__pycache__/ .cache/ diff --git a/__pycache__/embedlab.cpython-315.pyc b/__pycache__/embedlab.cpython-315.pyc deleted file mode 100644 index aaf7760..0000000 Binary files a/__pycache__/embedlab.cpython-315.pyc and /dev/null differ diff --git a/embedlab.py b/embedlab.py index 38c2608..79d8f04 100755 --- a/embedlab.py +++ b/embedlab.py @@ -2,6 +2,7 @@ print("Loading torch/transformers, please wait...", flush=True) +import atexit import os import readline @@ -22,6 +23,7 @@ MODELS = { PARAM_KEYS = ["pooling", "max_length", "batch_size"] DEFAULT_MODEL = "BAAI/bge-small-zh-v1.5" LAST_MODEL_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "last_model") +HISTORY_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".cache", "history") COMMANDS = ["/select", "/sim", "/tokens", "/params", "/set", "/reset", "/lsmod", "/help", "/quit", "/exit"] device = "cuda" if torch.cuda.is_available() else "cpu" @@ -202,6 +204,20 @@ def setup_readline(): readline.parse_and_bind("bind ^I rl_complete") else: readline.parse_and_bind("tab: complete") + try: + readline.read_history_file(HISTORY_FILE) + except OSError: + pass + readline.set_history_length(1000) + atexit.register(save_history) + + +def save_history(): + try: + os.makedirs(os.path.dirname(HISTORY_FILE), exist_ok=True) + readline.write_history_file(HISTORY_FILE) + except OSError: + pass HELP = """Commands: