Function test_real_llama

tests/test_llama.py:159–258

(llama_cpp_model_path)

def test_real_llama(llama_cpp_model_path):
    model = llama_cpp.Llama(
        llama_cpp_model_path,
        n_ctx=32,
        n_batch=32,
        n_ubatch=32,
        n_threads=multiprocessing.cpu_count(),
        n_threads_batch=multiprocessing.cpu_count(),
        logits_all=False,
        flash_attn=True,
    )

    output = model.create_completion(
        "The quick brown fox jumps over the lazy dog. The quick brown fox",
        max_tokens=6,
        top_k=50,
        top_p=0.9,
        temperature=0.0,
        seed=1337,
    )
    assert output["choices"][0]["text"] == " jumps over the lazy dog."

    output = model.create_completion(
        "The capital of france is paris, 'true' or 'false'?:\n",
        max_tokens=4,
        top_k=50,
        top_p=0.9,
        temperature=0.8,
        seed=1337,
        grammar=llama_cpp.LlamaGrammar.from_string("""
root ::= "true" | "false"
"""),
    )
    assert output["choices"][0]["text"] == "true"

    suffix = b"rot"
    tokens = model.tokenize(suffix, add_bos=True, special=True)

    def logit_processor_func(input_ids, logits):
        for token in tokens:
            logits[token] *= 1000
        return logits

    logit_processors = llama_cpp.LogitsProcessorList([logit_processor_func])

    output = model.create_completion(
        "The capital of france is par",
        max_tokens=4,
        top_k=50,
        top_p=0.9,
        temperature=0.8,
        seed=1337,
        logits_processor=logit_processors,
    )
    assert output["choices"][0]["text"].lower().startswith("rot")

    model.set_seed(1337)