diff --git a/.gitignore b/.gitignore index 21d0b89..fa326cf 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ .venv/ +__pycache__/ +.cache/ diff --git a/embedlab.py b/embedlab.py index 5dd0ca0..79d8f04 100755 --- a/embedlab.py +++ b/embedlab.py @@ -2,6 +2,8 @@ print("Loading torch/transformers, please wait...", flush=True) +import atexit +import os import readline import torch @@ -20,6 +22,8 @@ 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" @@ -38,6 +42,25 @@ def is_cached(name): ) +def load_last_model(): + """The model selected last time, or DEFAULT_MODEL if unknown.""" + try: + with open(LAST_MODEL_FILE) as f: + name = f.read().strip() + except OSError: + return DEFAULT_MODEL + return name if name in MODELS else DEFAULT_MODEL + + +def save_last_model(name): + try: + os.makedirs(os.path.dirname(LAST_MODEL_FILE), exist_ok=True) + with open(LAST_MODEL_FILE, "w") as f: + f.write(name) + except OSError: + pass + + def select_model(name): global current if name not in MODELS: @@ -49,6 +72,7 @@ def select_model(name): mod = AutoModel.from_pretrained(name).to(device).eval() _loaded[name] = (tok, mod) current = name + save_last_model(name) print(f"Current model: {name}") @@ -180,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: @@ -204,7 +242,7 @@ def main(): args = parser.parse_args() setup_readline() - select_model(DEFAULT_MODEL) + select_model(load_last_model()) print(HELP) while True: try: