Brunobkr/llama.cpp_AlgMor24_github

1import pytest2from openai import OpenAI3from utils import *4 5server: ServerProcess6 7@pytest.fixture(autouse=True)8def create_server():9    global server10    server = ServerPreset.tinyllama2()11 12 13@pytest.mark.parametrize(14    "model,system_prompt,user_prompt,max_tokens,re_content,n_prompt,n_predicted,finish_reason,jinja,chat_template",15    [16        (None, "Book", "Hey", 8, "But she couldn't", 69, 8, "length", False, None),17        (None, "Book", "Hey", 8, "But she couldn't", 69, 8, "length", True, None),18        (None, "Book", "What is the best book", 8, "(Suddenly)+|\\{ \" Sarax.", 77, 8, "length", False, None),19        (None, "Book", "What is the best book", 8, "(Suddenly)+|\\{ \" Sarax.", 77, 8, "length", True,  None),20        (None, "Book", "What is the best book", 8, "(Suddenly)+|\\{ \" Sarax.", 77, 8, "length", True, 'chatml'),21        (None, "Book", "What is the best book", 8, "^ blue",                    23, 8, "length", True, "This is not a chat template, it is"),22        ("codellama70b", "You are a coding assistant.", "Write the fibonacci function in c++.", 128, "(Aside|she|felter|alonger)+", 104, 128, "length", False, None),23        ("codellama70b", "You are a coding assistant.", "Write the fibonacci function in c++.", 128, "(Aside|she|felter|alonger)+", 104, 128, "length", True, None),24        (None, "Book", [{"type": "text", "text": "What is"}, {"type": "text", "text": "the best book"}], 8, "Whillicter", 79, 8, "length", False, None),25        (None, "Book", [{"type": "text", "text": "What is"}, {"type": "text", "text": "the best book"}], 8, "Whillicter", 79, 8, "length", True, None),26    ]27)28def test_chat_completion(model, system_prompt, user_prompt, max_tokens, re_content, n_prompt, n_predicted, finish_reason, jinja, chat_template):29    global server30    server.jinja = jinja31    server.chat_template = chat_template32    server.start()33    res = server.make_request("POST", "/chat/completions", data={34        "model": model,35        "max_tokens": max_tokens,36        "messages": [37            {"role": "system", "content": system_prompt},38            {"role": "user", "content": user_prompt},39        ],40    })41    assert res.status_code == 20042    assert "cmpl" in res.body["id"] # make sure the completion id has the expected format43    assert res.body["system_fingerprint"].startswith("b")44    # we no longer reflect back the model name, see https://github.com/ggml-org/llama.cpp/pull/1766845    # assert res.body["model"] == model if model is not None else server.model_alias46    assert res.body["usage"]["prompt_tokens"] == n_prompt47    assert res.body["usage"]["completion_tokens"] == n_predicted48    choice = res.body["choices"][0]49    assert "assistant" == choice["message"]["role"]50    assert match_regex(re_content, choice["message"]["content"]), f'Expected {re_content}, got {choice["message"]["content"]}'51    assert choice["finish_reason"] == finish_reason52 53 54def test_chat_completion_cached_tokens():55    global server56    server.n_slots = 157    server.start()58    seq = [59        ("1 2 3 4 5 6", 77, 0),60        ("1 2 3 4 5 6", 77, 76),61        ("1 2 3 4 5 9", 77, 51),62        ("1 2 3 9 9 9", 77, 47),63    ]64    for user_prompt, n_prompt, n_cache in seq:65        res = server.make_request("POST", "/chat/completions", data={66            "max_tokens": 8,67            "messages": [68                {"role": "system", "content": "Test"},69                {"role": "user", "content": user_prompt},70            ],71        })72        assert res.body["usage"]["prompt_tokens"] == n_prompt73        assert res.body["usage"]["prompt_tokens_details"]["cached_tokens"] == n_cache74 75@pytest.mark.parametrize(76    "system_prompt,user_prompt,max_tokens,re_content,n_prompt,n_predicted,finish_reason",77    [78        ("Book", "What is the best book", 8, "(Suddenly)+", 77, 8, "length"),79        ("You are a coding assistant.", "Write the fibonacci function in c++.", 128, "(Aside|she|felter|alonger)+", 104, 128, "length"),80    ]81)82def test_chat_completion_stream(system_prompt, user_prompt, max_tokens, re_content, n_prompt, n_predicted, finish_reason):83    global server84    server.model_alias = "llama-test-model"85    server.start()86    res = server.make_stream_request("POST", "/chat/completions", data={87        "max_tokens": max_tokens,88        "messages": [89            {"role": "system", "content": system_prompt},90            {"role": "user", "content": user_prompt},91        ],92        "stream": True,93    })94    content = ""95    last_cmpl_id = None96    for i, data in enumerate(res):97        if data["choices"]:98            choice = data["choices"][0]99            if i == 0:100                # Check first role message for stream=True101                assert choice["delta"]["content"] is None102                assert choice["delta"]["role"] == "assistant"103            else:104                assert "role" not in choice["delta"]105            assert data["system_fingerprint"].startswith("b")106            assert data["model"] == "llama-test-model"107            if last_cmpl_id is None:108                last_cmpl_id = data["id"]109            assert last_cmpl_id == data["id"] # make sure the completion id is the same for all events in the stream110            if choice["finish_reason"] in ["stop", "length"]:111                assert "content" not in choice["delta"]112                assert match_regex(re_content, content)113                assert choice["finish_reason"] == finish_reason114            else:115                assert choice["finish_reason"] is None116                content += choice["delta"]["content"] or ''117        else:118            assert data["usage"]["prompt_tokens"] == n_prompt119            assert data["usage"]["completion_tokens"] == n_predicted120 121 122def test_chat_completion_with_openai_library():123    global server124    server.start()125    client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")126    res = client.chat.completions.create(127        model="gpt-3.5-turbo-instruct",128        messages=[129            {"role": "system", "content": "Book"},130            {"role": "user", "content": "What is the best book"},131        ],132        max_tokens=8,133        seed=42,134        temperature=0.8,135    )136    assert res.system_fingerprint is not None and res.system_fingerprint.startswith("b")137    assert res.choices[0].finish_reason == "length"138    assert res.choices[0].message.content is not None139    assert match_regex("(Suddenly)+", res.choices[0].message.content)140 141 142def test_chat_template():143    global server144    server.chat_template = "llama3"145    server.debug = True  # to get the "__verbose" object in the response146    server.start()147    res = server.make_request("POST", "/chat/completions", data={148        "max_tokens": 8,149        "messages": [150            {"role": "system", "content": "Book"},151            {"role": "user", "content": "What is the best book"},152        ]153    })154    assert res.status_code == 200155    assert "__verbose" in res.body156    assert res.body["__verbose"]["prompt"] == "<s> <|start_header_id|>system<|end_header_id|>\n\nBook<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nWhat is the best book<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"157 158 159@pytest.mark.parametrize("prefill,re_prefill", [160    ("Whill", "Whill"),161    ([{"type": "text", "text": "Wh"}, {"type": "text", "text": "ill"}], "Wh\n\nill"),162])163def test_chat_template_assistant_prefill(prefill, re_prefill):164    global server165    server.jinja = True166    server.chat_template_file = "../../../models/templates/meta-llama-Llama-3.1-8B-Instruct.jinja"167    server.debug = True  # to get the "__verbose" object in the response168    server.start()169    res = server.make_request("POST", "/chat/completions", data={170        "max_tokens": 8,171        "messages": [172            {"role": "system", "content": "Book"},173            {"role": "user", "content": "What is the best book"},174            {"role": "assistant", "content": prefill},175        ]176    })177    assert res.status_code == 200178    assert "__verbose" in res.body179    assert res.body["__verbose"]["prompt"].endswith(f"<|start_header_id|>user<|end_header_id|>\n\nWhat is the best book<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n{re_prefill}")180 181 182def test_chat_template_continue_final_message_vllm_compat():183    """continue_final_message is the vLLM/transformers explicit alias for the prefill_assistant heuristic.184    Both must produce the same prompt."""185    global server186    server.jinja = True187    server.chat_template_file = "../../../models/templates/meta-llama-Llama-3.1-8B-Instruct.jinja"188    server.debug = True189    server.start()190    res = server.make_request("POST", "/chat/completions", data={191        "max_tokens": 8,192        "add_generation_prompt": False,193        "continue_final_message": True,194        "messages": [195            {"role": "system", "content": "Book"},196            {"role": "user", "content": "What is the best book"},197            {"role": "assistant", "content": "Whill"},198        ]199    })200    assert res.status_code == 200201    assert "__verbose" in res.body202    assert res.body["__verbose"]["prompt"].endswith("<|start_header_id|>user<|end_header_id|>\n\nWhat is the best book<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nWhill")203 204 205def test_chat_template_continue_final_message_mutual_exclusion():206    """add_generation_prompt and continue_final_message both set to true must be rejected"""207    global server208    server.chat_template = "llama3"209    server.start()210    res = server.make_request("POST", "/chat/completions", data={211        "max_tokens": 8,212        "add_generation_prompt": True,213        "continue_final_message": True,214        "messages": [215            {"role": "user", "content": "Hi"},216            {"role": "assistant", "content": "Hello"},217        ]218    })219    assert res.status_code == 400220 221 222def test_apply_chat_template():223    global server224    server.chat_template = "command-r"225    server.start()226    res = server.make_request("POST", "/apply-template", data={227        "messages": [228            {"role": "system", "content": "You are a test."},229            {"role": "user", "content":"Hi there"},230        ]231    })232    assert res.status_code == 200233    assert "prompt" in res.body234    assert res.body["prompt"] == "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>You are a test.<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Hi there<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"235 236 237@pytest.mark.parametrize("response_format,n_predicted,re_content", [238    ({"type": "json_object", "schema": {"const": "42"}}, 6, "\"42\""),239    ({"type": "json_object", "schema": {"items": [{"type": "integer"}]}}, 10, "[ -3000 ]"),240    ({"type": "json_schema", "json_schema": {"schema": {"const": "foooooo"}}}, 10, "\"foooooo\""),241    ({"type": "json_object"}, 10, "(\\{|John)+"),242    ({"type": "sound"}, 0, None),243    # invalid response format (expected to fail)244    ({"type": "json_object", "schema": 123}, 0, None),245    ({"type": "json_object", "schema": {"type": 123}}, 0, None),246    ({"type": "json_object", "schema": {"type": "hiccup"}}, 0, None),247])248def test_completion_with_response_format(response_format: dict, n_predicted: int, re_content: str | None):249    global server250    server.start()251    res = server.make_request("POST", "/chat/completions", data={252        "max_tokens": n_predicted,253        "messages": [254            {"role": "system", "content": "You are a coding assistant."},255            {"role": "user", "content": "Write an example"},256        ],257        "response_format": response_format,258    })259    if re_content is not None:260        assert res.status_code == 200261        choice = res.body["choices"][0]262        assert match_regex(re_content, choice["message"]["content"])263    else:264        assert res.status_code == 400265        assert "error" in res.body266 267 268@pytest.mark.parametrize("jinja,json_schema,n_predicted,re_content", [269    (False, {"const": "42"}, 6, "\"42\""),270    (True, {"const": "42"}, 6, "\"42\""),271])272def test_completion_with_json_schema(jinja: bool, json_schema: dict, n_predicted: int, re_content: str):273    global server274    server.jinja = jinja275    server.debug = True276    server.start()277    res = server.make_request("POST", "/chat/completions", data={278        "max_tokens": n_predicted,279        "messages": [280            {"role": "system", "content": "You are a coding assistant."},281            {"role": "user", "content": "Write an example"},282        ],283        "json_schema": json_schema,284    })285    assert res.status_code == 200, f'Expected 200, got {res.status_code}'286    choice = res.body["choices"][0]287    assert match_regex(re_content, choice["message"]["content"]), f'Expected {re_content}, got {choice["message"]["content"]}'288 289 290@pytest.mark.parametrize("jinja,grammar,n_predicted,re_content", [291    (False, 'root ::= "a"{5,5}', 6, "a{5,5}"),292    (True, 'root ::= "a"{5,5}', 6, "a{5,5}"),293])294def test_completion_with_grammar(jinja: bool, grammar: str, n_predicted: int, re_content: str):295    global server296    server.jinja = jinja297    server.start()298    res = server.make_request("POST", "/chat/completions", data={299        "max_tokens": n_predicted,300        "messages": [301            {"role": "user", "content": "Does not matter what I say, does it?"},302        ],303        "grammar": grammar,304    })305    assert res.status_code == 200, res.body306    choice = res.body["choices"][0]307    assert match_regex(re_content, choice["message"]["content"]), choice["message"]["content"]308 309 310def test_completion_with_invalid_grammar():311    global server312    server.start()313    res = server.make_request("POST", "/chat/completions", data={314        "max_tokens": 8,315        "messages": [316            {"role": "user", "content": "Does not matter what I say, does it?"},317        ],318        "grammar": "root ::= this is (not valid GBNF",319    })320    assert res.status_code == 400, res.body321    assert "error" in res.body322 323 324@pytest.mark.parametrize("messages", [325    None,326    "string",327    [123],328    [{}],329    [{"role": 123}],330    [{"role": "system", "content": 123}],331    # [{"content": "hello"}], # TODO: should not be a valid case332    [{"role": "system", "content": "test"}, {}],333    [{"role": "user", "content": "test"}, {"role": "assistant", "content": "test"}, {"role": "assistant", "content": "test"}],334])335def test_invalid_chat_completion_req(messages):336    global server337    server.start()338    res = server.make_request("POST", "/chat/completions", data={339        "messages": messages,340    })341    assert res.status_code == 400 or res.status_code == 500342    assert "error" in res.body343 344 345def test_chat_completion_with_timings_per_token():346    global server347    server.start()348    res = server.make_stream_request("POST", "/chat/completions", data={349        "max_tokens": 10,350        "messages": [{"role": "user", "content": "test"}],351        "stream": True,352        "stream_options": {"include_usage": True},353        "timings_per_token": True,354    })355    stats_received = False356    for i, data in enumerate(res):357        if i == 0:358            # Check first role message for stream=True359            assert data["choices"][0]["delta"]["content"] is None360            assert data["choices"][0]["delta"]["role"] == "assistant"361            assert "timings" not in data, f'First event should not have timings: {data}'362        else:363            if data["choices"]:364                assert "role" not in data["choices"][0]["delta"]365            else:366                assert "timings" in data367                assert "prompt_per_second" in data["timings"]368                assert "predicted_per_second" in data["timings"]369                assert "predicted_n" in data["timings"]370                assert data["timings"]["predicted_n"] <= 10371                stats_received = True372    assert stats_received373 374 375def test_logprobs():376    global server377    server.start()378    client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")379    res = client.chat.completions.create(380        model="gpt-3.5-turbo-instruct",381        temperature=0.0,382        messages=[383            {"role": "system", "content": "Book"},384            {"role": "user", "content": "What is the best book"},385        ],386        max_tokens=5,387        logprobs=True,388        top_logprobs=10,389    )390    output_text = res.choices[0].message.content391    aggregated_text = ''392    assert res.choices[0].logprobs is not None393    assert res.choices[0].logprobs.content is not None394    for token in res.choices[0].logprobs.content:395        aggregated_text += token.token396        assert token.logprob <= 0.0397        assert token.bytes is not None398        assert len(token.top_logprobs) > 0399    assert aggregated_text == output_text400 401 402def test_logprobs_stream():403    global server404    server.start()405    client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")406    res = client.chat.completions.create(407        model="gpt-3.5-turbo-instruct",408        temperature=0.0,409        messages=[410            {"role": "system", "content": "Book"},411            {"role": "user", "content": "What is the best book"},412        ],413        max_tokens=5,414        logprobs=True,415        top_logprobs=10,416        stream=True,417    )418    output_text = ''419    aggregated_text = ''420    for i, data in enumerate(res):421        if data.choices:422            choice = data.choices[0]423            if i == 0:424                # Check first role message for stream=True425                assert choice.delta.content is None426                assert choice.delta.role == "assistant"427            else:428                assert choice.delta.role is None429                if choice.finish_reason is None:430                    if choice.delta.content:431                        output_text += choice.delta.content432                    assert choice.logprobs is not None433                    assert choice.logprobs.content is not None434                    for token in choice.logprobs.content:435                        aggregated_text += token.token436                        assert token.logprob <= 0.0437                        assert token.bytes is not None438                        assert token.top_logprobs is not None439                        assert len(token.top_logprobs) > 0440    assert aggregated_text == output_text441 442 443def test_logit_bias():444    global server445    server.start()446 447    exclude = ["i", "I", "the", "The", "to", "a", "an", "be", "is", "was", "but", "But", "and", "And", "so", "So", "you", "You", "he", "He", "she", "She", "we", "We", "they", "They", "it", "It", "his", "His", "her", "Her", "book", "Book"]448 449    res = server.make_request("POST", "/tokenize", data={450        "content": " " + " ".join(exclude) + " ",451    })452    assert res.status_code == 200453    tokens = res.body["tokens"]454    logit_bias = {tok: -100 for tok in tokens}455 456    client = OpenAI(api_key="dummy", base_url=f"http://{server.server_host}:{server.server_port}/v1")457    res = client.chat.completions.create(458        model="gpt-3.5-turbo-instruct",459        temperature=0.0,460        messages=[461            {"role": "system", "content": "Book"},462            {"role": "user", "content": "What is the best book"},463        ],464        max_tokens=64,465        logit_bias=logit_bias466    )467    output_text = res.choices[0].message.content468    assert output_text469    assert all(output_text.find(" " + tok + " ") == -1 for tok in exclude)470 471def test_context_size_exceeded():472    global server473    server.start()474    res = server.make_request("POST", "/chat/completions", data={475        "messages": [476            {"role": "system", "content": "Book"},477            {"role": "user", "content": "What is the best book"},478        ] * 100, # make the prompt too long479    })480    assert res.status_code == 400481    assert "error" in res.body482    assert res.body["error"]["type"] == "exceed_context_size_error"483    assert res.body["error"]["n_prompt_tokens"] > 0484    assert server.n_ctx is not None485    assert server.n_slots is not None486    assert res.body["error"]["n_ctx"] == server.n_ctx // server.n_slots487 488 489def test_context_size_exceeded_stream():490    global server491    server.start()492    try:493        for _ in server.make_stream_request("POST", "/chat/completions", data={494            "messages": [495                {"role": "system", "content": "Book"},496                {"role": "user", "content": "What is the best book"},497            ] * 100, # make the prompt too long498            "stream": True}):499                pass500        assert False, "Should have failed"501    except ServerError as e:502        assert e.code == 400503        assert "error" in e.body504        assert e.body["error"]["type"] == "exceed_context_size_error"505        assert e.body["error"]["n_prompt_tokens"] > 0506        assert server.n_ctx is not None507        assert server.n_slots is not None508        assert e.body["error"]["n_ctx"] == server.n_ctx // server.n_slots509 510 511@pytest.mark.parametrize(512    "n_batch,batch_count,reuse_cache",513    [514        (64, 4, False),515        (64, 2, True),516    ]517)518def test_return_progress(n_batch, batch_count, reuse_cache):519    global server520    server.n_batch = n_batch521    server.n_ctx = 256522    server.n_slots = 1523    server.start()524    def make_cmpl_request():525        return server.make_stream_request("POST", "/chat/completions", data={526            "max_tokens": 10,527            "messages": [528                {"role": "user", "content": "This is a test" * 10},529            ],530            "stream": True,531            "return_progress": True,532        })533    if reuse_cache:534        # make a first request to populate the cache535        res0 = make_cmpl_request()536        for _ in res0:537            pass # discard the output538 539    res = make_cmpl_request()540    last_progress = None541    total_batch_count = 0542 543    for data in res:544        cur_progress = data.get("prompt_progress", None)545        if cur_progress is None:546            continue547        if total_batch_count == 0:548            # first progress report must have n_cache == n_processed549            assert cur_progress["total"] > 0550            assert cur_progress["cache"] == cur_progress["processed"]551            if reuse_cache:552                # when reusing cache, we expect some cached tokens553                assert cur_progress["cache"] > 0554        if last_progress is not None:555            assert cur_progress["total"] == last_progress["total"]556            assert cur_progress["cache"] == last_progress["cache"]557            assert cur_progress["processed"] > last_progress["processed"]558        total_batch_count += 1559        last_progress = cur_progress560 561    # last progress should indicate completion (all tokens processed)562    assert last_progress is not None563    assert last_progress["total"] > 0564    assert last_progress["processed"] == last_progress["total"]565    assert total_batch_count == batch_count566 567 568def test_chat_completions_multiple_choices():569    global server570    server.start()571    # make sure cache can be reused across multiple choices and multiple requests572    # ref: https://github.com/ggml-org/llama.cpp/pull/18663573    for _ in range(2):574        res = server.make_request("POST", "/chat/completions", data={575            "max_tokens": 8,576            "n": 2,577            "messages": [578                {"role": "system", "content": "Book"},579                {"role": "user", "content": "What is the best book"},580            ],581            # test forcing the same slot to be used582            # the scheduler should not be locked up in this case583            "id_slot": 0,584        })585        assert res.status_code == 200586        assert len(res.body["choices"]) == 2587        for choice in res.body["choices"]:588            assert "assistant" == choice["message"]["role"]589            assert choice["finish_reason"] == "length"590 591 592def test_chat_completions_token_count():593    global server594    server.start()595    # make sure cache can be reused across multiple choices and multiple requests596    # ref: https://github.com/ggml-org/llama.cpp/pull/18663597    for _ in range(2):598        res = server.make_request("POST", "/chat/completions/input_tokens", data={599            "messages": [600                {"role": "system", "content": "Book"},601                {"role": "user", "content": "What is the best book"},602            ],603        })604        assert res.status_code == 200605        assert res.body["input_tokens"] > 5606 607 608def test_verbose_debug():609    global server610    server.start()611    for verbose in [True, False]:612        res = server.make_request("POST", "/chat/completions", data={613            "max_tokens": 2,614            "messages": [615                {"role": "system", "content": "Book"},616                {"role": "user", "content": "What is the best book"},617            ],618            "verbose": verbose,619        })620        assert res.status_code == 200621        if verbose:622            assert "__verbose" in res.body623            assert "Book" in res.body["__verbose"]["prompt"]624        else:625            assert "__verbose" not in res.body626