KBaba7/llama.cpp

1import pytest2from utils import *3 4# We use a F16 MOE gguf as main model, and q4_0 as draft model5 6server = ServerPreset.stories15m_moe()7 8MODEL_DRAFT_FILE_URL = "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf"9 10def create_server():11    global server12    server = ServerPreset.stories15m_moe()13    # set default values14    server.model_draft = download_file(MODEL_DRAFT_FILE_URL)15    server.draft_min = 416    server.draft_max = 817 18 19@pytest.fixture(scope="module", autouse=True)20def fixture_create_server():21    return create_server()22 23 24def test_with_and_without_draft():25    global server26    server.model_draft = None  # disable draft model27    server.start()28    res = server.make_request("POST", "/completion", data={29        "prompt": "I believe the meaning of life is",30        "temperature": 0.0,31        "top_k": 1,32    })33    assert res.status_code == 20034    content_no_draft = res.body["content"]35    server.stop()36 37    # create new server with draft model38    create_server()39    server.start()40    res = server.make_request("POST", "/completion", data={41        "prompt": "I believe the meaning of life is",42        "temperature": 0.0,43        "top_k": 1,44    })45    assert res.status_code == 20046    content_draft = res.body["content"]47 48    assert content_no_draft == content_draft49 50 51def test_different_draft_min_draft_max():52    global server53    test_values = [54        (1, 2),55        (1, 4),56        (4, 8),57        (4, 12),58        (8, 16),59    ]60    last_content = None61    for draft_min, draft_max in test_values:62        server.stop()63        server.draft_min = draft_min64        server.draft_max = draft_max65        server.start()66        res = server.make_request("POST", "/completion", data={67            "prompt": "I believe the meaning of life is",68            "temperature": 0.0,69            "top_k": 1,70        })71        assert res.status_code == 20072        if last_content is not None:73            assert last_content == res.body["content"]74        last_content = res.body["content"]75 76 77def test_slot_ctx_not_exceeded():78    global server79    server.n_ctx = 6480    server.start()81    res = server.make_request("POST", "/completion", data={82        "prompt": "Hello " * 56,83        "temperature": 0.0,84        "top_k": 1,85        "speculative.p_min": 0.0,86    })87    assert res.status_code == 20088    assert len(res.body["content"]) > 089 90 91def test_with_ctx_shift():92    global server93    server.n_ctx = 6494    server.start()95    res = server.make_request("POST", "/completion", data={96        "prompt": "Hello " * 56,97        "temperature": 0.0,98        "top_k": 1,99        "n_predict": 64,100        "speculative.p_min": 0.0,101    })102    assert res.status_code == 200103    assert len(res.body["content"]) > 0104    assert res.body["tokens_predicted"] == 64105    assert res.body["truncated"] == True106 107 108@pytest.mark.parametrize("n_slots,n_requests", [109    (1, 2),110    (2, 2),111])112def test_multi_requests_parallel(n_slots: int, n_requests: int):113    global server114    server.n_slots = n_slots115    server.start()116    tasks = []117    for _ in range(n_requests):118        tasks.append((server.make_request, ("POST", "/completion", {119            "prompt": "I believe the meaning of life is",120            "temperature": 0.0,121            "top_k": 1,122        })))123    results = parallel_function_calls(tasks)124    for res in results:125        assert res.status_code == 200126        assert match_regex("(wise|kind|owl|answer)+", res.body["content"])127