diff --git a/.gitignore b/.gitignore index 21d0b89..0fea4f7 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ .venv/ +.cache/ diff --git a/__pycache__/embedlab.cpython-315.pyc b/__pycache__/embedlab.cpython-315.pyc new file mode 100644 index 0000000..aaf7760 Binary files /dev/null and b/__pycache__/embedlab.cpython-315.pyc differ diff --git a/embedlab.py b/embedlab.py index 5dd0ca0..38c2608 100755 --- a/embedlab.py +++ b/embedlab.py @@ -2,6 +2,7 @@ print("Loading torch/transformers, please wait...", flush=True) +import os import readline import torch @@ -20,6 +21,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") COMMANDS = ["/select", "/sim", "/tokens", "/params", "/set", "/reset", "/lsmod", "/help", "/quit", "/exit"] device = "cuda" if torch.cuda.is_available() else "cpu" @@ -38,6 +40,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 +70,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}") @@ -204,7 +226,7 @@ def main(): args = parser.parse_args() setup_readline() - select_model(DEFAULT_MODEL) + select_model(load_last_model()) print(HELP) while True: try: