Brunobkr/llama.cpp_AlgMor24_github

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/tiny-llamas/resolve/main/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.spec_type = "draft-simple"16    server.spec_draft_n_min = 417    server.spec_draft_n_max = 818    server.fa = "off"19 20 21@pytest.fixture(autouse=True)22def fixture_create_server():23    return create_server()24 25 26def test_with_and_without_draft():27    global server28    request = {29        "prompt": "I believe the meaning of life is",30        "temperature": 0.8,31        "top_k": 40,32        "seed": 4242,33        "n_predict": 16,34        "return_tokens": True,35    }36 37    server.model_draft = None  # disable draft model38    server.spec_type = None39    server.backend_sampling = True40    server.start()41    res = server.make_request("POST", "/completion", data=request)42    assert res.status_code == 20043    tokens_no_draft = res.body["tokens"]44    server.stop()45 46    # create new server with draft model47    create_server()48    server.backend_sampling = True49    server.start()50    res = server.make_request("POST", "/completion", data=request)51    assert res.status_code == 20052    assert res.body["timings"]["draft_n"] > 053    tokens_draft = res.body["tokens"]54 55    assert tokens_no_draft == tokens_draft56 57 58def test_different_draft_min_draft_max():59    global server60    test_values = [61        (1, 2),62        (1, 4),63        (4, 8),64        (4, 12),65        (8, 16),66    ]67    last_content = None68    for draft_min, draft_max in test_values:69        server.stop()70        server.spec_draft_n_min = draft_min71        server.spec_draft_n_max = draft_max72        server.start()73        res = server.make_request("POST", "/completion", data={74            "prompt": "I believe the meaning of life is",75            "temperature": 0.0,76            "top_k": 1,77            "n_predict": 16,78        })79        assert res.status_code == 20080        if last_content is not None:81            assert last_content == res.body["content"]82        last_content = res.body["content"]83 84 85def test_slot_ctx_not_exceeded():86    global server87    server.n_ctx = 25688    server.start()89    res = server.make_request("POST", "/completion", data={90        "prompt": "Hello " * 248,91        "temperature": 0.0,92        "top_k": 1,93        "speculative.p_min": 0.0,94    })95    assert res.status_code == 20096    assert len(res.body["content"]) > 097 98 99def test_with_ctx_shift():100    global server101    server.n_ctx = 256102    server.enable_ctx_shift = True103    server.start()104    res = server.make_request("POST", "/completion", data={105        "prompt": "Hello " * 248,106        "temperature": 0.0,107        "top_k": 1,108        "n_predict": 256,109        "speculative.p_min": 0.0,110    })111    assert res.status_code == 200112    assert len(res.body["content"]) > 0113    assert res.body["tokens_predicted"] == 256114    assert res.body["truncated"] == True115 116 117@pytest.mark.parametrize("n_slots,n_requests", [118    (1, 2),119    (2, 2),120])121def test_multi_requests_parallel(n_slots: int, n_requests: int):122    global server123    server.n_slots = n_slots124    server.start()125    tasks = []126    for _ in range(n_requests):127        tasks.append((server.make_request, ("POST", "/completion", {128            "prompt": "I believe the meaning of life is",129            "temperature": 0.0,130            "top_k": 1,131        })))132    results = parallel_function_calls(tasks)133    for res in results:134        assert res.status_code == 200135        assert match_regex("(wise|kind|owl|answer)+", res.body["content"])136