diff --git a/src/whichllm/cli.py b/src/whichllm/cli.py index 46b1b71..dafea1f 100644 --- a/src/whichllm/cli.py +++ b/src/whichllm/cli.py @@ -1133,12 +1133,15 @@ def _generate_chat_script(model, variant, context_length: int, cpu_only: bool) - """Generate a self-contained Python chat script for any model type.""" if variant: n_gpu = 0 if cpu_only else -1 - return f'''\ + return f"""\ from huggingface_hub import hf_hub_download from llama_cpp import Llama -print("Downloading {model.id} ({variant.quant_type})...") -model_path = hf_hub_download(repo_id="{model.id}", filename="{variant.filename}") +model_id = {model.id!r} +filename = {variant.filename!r} +quant_type = {variant.quant_type!r} +print(f"Downloading {{model_id}} ({{quant_type}})...") +model_path = hf_hub_download(repo_id=model_id, filename=filename) print("Loading model...") llm = Llama( model_path=model_path, @@ -1169,18 +1172,18 @@ def _generate_chat_script(model, variant, context_length: int, cpu_only: bool) - print() messages.append({{"role": "assistant", "content": full}}) print("\\nBye!") -''' +""" device_map = '"cpu"' if cpu_only else '"auto"' dtype = "torch.float32" if cpu_only else '"auto"' - return f'''\ + return f"""\ import shutil import tempfile import torch from threading import Thread from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer -model_id = "{model.id}" +model_id = {model.id!r} offload_folder = tempfile.mkdtemp(prefix="whichllm_transformers_offload_") try: print(f"Loading {{model_id}}...") @@ -1232,7 +1235,7 @@ def _generate_chat_script(model, variant, context_length: int, cpu_only: bool) - except NameError: pass shutil.rmtree(offload_folder, ignore_errors=True) -''' +""" @app.command() @@ -1414,12 +1417,12 @@ def snippet( deps, _ = _resolve_model_deps(model, variant) if variant: - code = f'''\ + code = f"""\ from llama_cpp import Llama llm = Llama.from_pretrained( - repo_id="{model.id}", - filename="{variant.filename}", + repo_id={model.id!r}, + filename={variant.filename!r}, n_ctx=4096, n_gpu_layers=-1, # -1 = all layers on GPU, 0 = CPU only verbose=False, @@ -1429,12 +1432,12 @@ def snippet( messages=[{{"role": "user", "content": "Hello!"}}], ) print(output["choices"][0]["message"]["content"]) -''' +""" else: - code = f'''\ + code = f"""\ from transformers import AutoModelForCausalLM, AutoTokenizer -model_id = "{model.id}" +model_id = {model.id!r} tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( model_id, device_map="auto", torch_dtype="auto", trust_remote_code=True, @@ -1443,7 +1446,7 @@ def snippet( inputs = tokenizer("Hello!", return_tensors="pt").to(model.device) outputs = model.generate(**inputs, max_new_tokens=256) print(tokenizer.decode(outputs[0], skip_special_tokens=True)) -''' +""" dep_str = " ".join(f"--with {d}" for d in deps) console.print(f"\n[bold]{model.id}[/]") diff --git a/tests/test_cli.py b/tests/test_cli.py index 5ebf8fb..c3a7eb8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,5 +1,7 @@ """Tests for CLI helper logic.""" +import ast + import httpx import pytest from typer import Exit @@ -1175,6 +1177,53 @@ def test_transformers_chat_script_provides_disk_offload_folder(): assert "shutil.rmtree(offload_folder, ignore_errors=True)" in script +def test_gguf_chat_script_treats_hf_metadata_as_literals(): + model = _make_model(model_id="org/Test-7B") + filename = 'weights-Q4_K_M.gguf"); print("injected"); #.gguf' + variant = GGUFVariant( + filename=filename, + quant_type="Q4_K_M", + file_size_bytes=1, + ) + + tree = ast.parse( + _generate_chat_script(model, variant, context_length=4096, cpu_only=False) + ) + assignments = { + node.targets[0].id: ast.literal_eval(node.value) + for node in tree.body + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id in {"model_id", "filename", "quant_type"} + } + + assert assignments == { + "model_id": model.id, + "filename": filename, + "quant_type": variant.quant_type, + } + + +def test_snippet_treats_hf_metadata_as_literals(monkeypatch): + filename = 'weights-Q4_K_M.gguf"); print("injected"); #.gguf' + model = _make_model(model_id="org/Test-7B") + model.gguf_variants = [ + GGUFVariant( + filename=filename, + quant_type="Q4_K_M", + file_size_bytes=1, + ) + ] + monkeypatch.setattr(cli_mod, "_load_models", lambda refresh: [model]) + + result = CliRunner().invoke(app, ["snippet", "Test-7B"]) + + assert result.exit_code == 0 + assert f"repo_id={model.id!r}" in result.stdout + assert f"filename={filename!r}" in result.stdout + + def test_run_auto_pick_resolves_ranked_gguf_before_launch(monkeypatch): selected = ModelInfo( id="Qwen/Qwen3.6-27B",