KBaba7/llama.cpp

1import pytest2from utils import *3 4server = ServerPreset.jina_reranker_tiny()5 6 7@pytest.fixture(scope="module", autouse=True)8def create_server():9    global server10    server = ServerPreset.jina_reranker_tiny()11 12 13def test_rerank():14    global server15    server.start()16    res = server.make_request("POST", "/rerank", data={17        "query": "Machine learning is",18        "documents": [19            "A machine is a physical system that uses power to apply forces and control movement to perform an action. The term is commonly applied to artificial devices, such as those employing engines or motors, but also to natural biological macromolecules, such as molecular machines.",20            "Learning is the process of acquiring new understanding, knowledge, behaviors, skills, values, attitudes, and preferences. The ability to learn is possessed by humans, non-human animals, and some machines; there is also evidence for some kind of learning in certain plants.",21            "Machine learning is a field of study in artificial intelligence concerned with the development and study of statistical algorithms that can learn from data and generalize to unseen data, and thus perform tasks without explicit instructions.",22            "Paris, capitale de la France, est une grande ville européenne et un centre mondial de l'art, de la mode, de la gastronomie et de la culture. Son paysage urbain du XIXe siècle est traversé par de larges boulevards et la Seine."23        ]24    })25    assert res.status_code == 20026    assert len(res.body["results"]) == 427 28    most_relevant = res.body["results"][0]29    least_relevant = res.body["results"][0]30    for doc in res.body["results"]:31        if doc["relevance_score"] > most_relevant["relevance_score"]:32            most_relevant = doc33        if doc["relevance_score"] < least_relevant["relevance_score"]:34            least_relevant = doc35 36    assert most_relevant["relevance_score"] > least_relevant["relevance_score"]37    assert most_relevant["index"] == 238    assert least_relevant["index"] == 339 40 41@pytest.mark.parametrize("documents", [42    [],43    None,44    123,45    [1, 2, 3],46])47def test_invalid_rerank_req(documents):48    global server49    server.start()50    res = server.make_request("POST", "/rerank", data={51        "query": "Machine learning is",52        "documents": documents,53    })54    assert res.status_code == 40055    assert "error" in res.body56 57 58@pytest.mark.parametrize(59    "query,doc1,doc2,n_tokens",60    [61        ("Machine learning is", "A machine", "Learning is", 19),62        ("Which city?", "Machine learning is ", "Paris, capitale de la", 26),63    ]64)65def test_rerank_usage(query, doc1, doc2, n_tokens):66    global server67    server.start()68 69    res = server.make_request("POST", "/rerank", data={70        "query": query,71        "documents": [72            doc1,73            doc2,74        ]75    })76    assert res.status_code == 20077    assert res.body['usage']['prompt_tokens'] == res.body['usage']['total_tokens']78    assert res.body['usage']['prompt_tokens'] == n_tokens79