Function test_real_model

tests/test_llama.py:100–156

(llama_cpp_model_path)

def test_real_model(llama_cpp_model_path):
    import os

    assert os.path.exists(llama_cpp_model_path)

    params = llama_cpp.llama_model_default_params()
    params.use_mmap = llama_cpp.llama_supports_mmap()
    params.use_mlock = llama_cpp.llama_supports_mlock()
    params.check_tensors = False

    model = internals.LlamaModel(path_model=llama_cpp_model_path, params=params)

    cparams = llama_cpp.llama_context_default_params()
    cparams.n_ctx = 16
    cparams.n_batch = 16
    cparams.n_ubatch = 16
    cparams.n_threads = multiprocessing.cpu_count()
    cparams.n_threads_batch = multiprocessing.cpu_count()
    cparams.logits_all = False
    cparams.flash_attn_type = llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED

    context = internals.LlamaContext(model=model, params=cparams)
    tokens = model.tokenize(b"Hello, world!", add_bos=True, special=True)

    assert tokens == [9419, 11, 1814, 0]

    tokens = model.tokenize(
        b"The quick brown fox jumps over the lazy dog. The quick brown fox jumps ",
        add_bos=True,
        special=True,
    )
    prompt_token_count = len(tokens)

    batch = internals.LlamaBatch(n_tokens=len(tokens), embd=0, n_seq_max=1)

    seed = 1337
    sampler = internals.LlamaSampler()
    sampler.add_top_k(50)
    sampler.add_top_p(0.9, 1)
    sampler.add_temp(0.8)
    sampler.add_dist(seed)

    result = tokens
    n_eval = 0
    for _ in range(4):
        batch.set_batch(tokens, n_past=n_eval, logits_all=False)
        context.decode(batch)
        n_eval += len(tokens)
        token_id = sampler.sample(context, -1)
        tokens = [token_id]
        result += tokens

    output = result[prompt_token_count:]
    output_text = model.detokenize(output, special=True)
    # Low-level sampling output varies across CPU and Metal backends.
    assert len(output) == 4
    assert output_text