test_gritlm.py 7.58 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import numpy as np
4
5
6
7
import openai
import pytest
from scipy.spatial.distance import cosine

8
from vllm import LLM, SamplingParams
9
from vllm.config import ModelConfig, RendererConfig
10
11
12
13
14

from ....utils import RemoteOpenAIServer

MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
MAX_MODEL_LEN = 4000
15
ATOL = 0.002
16
17
18
19
20
21


def _arr(arr):
    """
    Convert a list of integers to an array of integers.
    """
22
    return np.array(arr)
23
24


25
def test_find_array():
26
    from vllm.model_executor.models.gritlm import GritLMMeanPool
27

28
29
    model_config = ModelConfig(
        MODEL_NAME,
30
        runner="pooling",
31
32
33
        dtype="bfloat16",
        seed=0,
    )
34
35
    renderer_config = RendererConfig(model_config=model_config)
    pooling = GritLMMeanPool(renderer_config=renderer_config)
36

37
    arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
38

39
40
41
42
43
44
    assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3
    assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3
    assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1
    assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=3) == -1
    assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=4) == 3
    assert pooling._find_array(arr, _arr([3, 5]), start_idx=0) == -1
45

46
    with pytest.raises(ValueError):
47
        pooling._find_array(arr, _arr([3, 4, 5]), start_idx=-1)
48
49


50
def run_llm_encode(
51
    llm: LLM,
52
53
    queries: list[str],
    instruction: str,
54
55
) -> list[list[float]]:
    outputs = llm.embed([instruction + q for q in queries])
56
57
58
    return [output.outputs.embedding for output in outputs]


59
async def run_client_embeddings(
60
    client: openai.AsyncOpenAI,
61
62
    queries: list[str],
    instruction: str,
63
) -> list[list[float]]:
64
65
66
67
68
69
70
71
    outputs = await client.embeddings.create(
        model=MODEL_NAME,
        input=[instruction + q for q in queries],
    )
    return [data.embedding for data in outputs.data]


def gritlm_instruction(instruction):
72
73
74
    return (
        "<|user|>\n" + instruction + "\n<|embed|>\n" if instruction else "<|embed|>\n"
    )
75
76
77
78
79
80
81
82


def get_test_data():
    """
    Grabbed this test data and the expected values from
    README.md in https://github.com/ContextualAI/gritlm
    """
    q_instruction = gritlm_instruction(
83
84
        "Given a scientific paper title, retrieve the paper's abstract",
    )
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    queries = [
        "Bitcoin: A Peer-to-Peer Electronic Cash System",
        "Generative Representational Instruction Tuning",
    ]

    d_instruction = gritlm_instruction("")
    documents = [
        # ruff: noqa: E501
        "A purely peer-to-peer version of electronic cash would allow online payments to be sent directly from one party to another without going through a financial institution. Digital signatures provide part of the solution, but the main benefits are lost if a trusted third party is still required to prevent double-spending. We propose a solution to the double-spending problem using a peer-to-peer network. The network timestamps transactions by hashing them into an ongoing chain of hash-based proof-of-work, forming a record that cannot be changed without redoing the proof-of-work. The longest chain not only serves as proof of the sequence of events witnessed, but proof that it came from the largest pool of CPU power. As long as a majority of CPU power is controlled by nodes that are not cooperating to attack the network, they'll generate the longest chain and outpace attackers. The network itself requires minimal structure. Messages are broadcast on a best effort basis, and nodes can leave and rejoin the network at will, accepting the longest proof-of-work chain as proof of what happened while they were gone.",
        "All text-based language problems can be reduced to either generation or embedding. Current models only perform well at one or the other. We introduce generative representational instruction tuning (GRIT) whereby a large language model is trained to handle both generative and embedding tasks by distinguishing between them through instructions. Compared to other open models, our resulting GritLM 7B sets a new state of the art on the Massive Text Embedding Benchmark (MTEB) and outperforms all models up to its size on a range of generative tasks. By scaling up further, GritLM 8X7B outperforms all open generative language models that we tried while still being among the best embedding models. Notably, we find that GRIT matches training on only generative or embedding data, thus we can unify both at no performance loss. Among other benefits, the unification via GRIT speeds up Retrieval-Augmented Generation (RAG) by > 60% for long documents, by no longer requiring separate retrieval and generation models. Models, code, etc. are freely available at https://github.com/ContextualAI/gritlm.",
    ]

    return queries, q_instruction, documents, d_instruction


100
def validate_embed_output(q_rep: list[list[float]], d_rep: list[list[float]]):
101
    cosine_sim_q0_d0 = 1 - cosine(q_rep[0], d_rep[0])
102
    assert cosine_sim_q0_d0 == pytest.approx(0.609, abs=ATOL)
103
104

    cosine_sim_q0_d1 = 1 - cosine(q_rep[0], d_rep[1])
105
    assert cosine_sim_q0_d1 == pytest.approx(0.101, abs=ATOL)
106
107

    cosine_sim_q1_d0 = 1 - cosine(q_rep[1], d_rep[0])
108
    assert cosine_sim_q1_d0 == pytest.approx(0.120, abs=ATOL)
109
110

    cosine_sim_q1_d1 = 1 - cosine(q_rep[1], d_rep[1])
111
    assert cosine_sim_q1_d1 == pytest.approx(0.534, abs=ATOL)
112
113


114
115
def test_gritlm_offline_embedding(vllm_runner):
    queries, q_instruction, documents, d_instruction = get_test_data()
116

117
    with vllm_runner(
118
119
120
        MODEL_NAME,
        runner="pooling",
        max_model_len=MAX_MODEL_LEN,
121
    ) as vllm_model:
122
        llm = vllm_model.llm
123

124
125
126
127
128
129
130
131
132
133
        d_rep = run_llm_encode(
            llm,
            documents,
            d_instruction,
        )
        q_rep = run_llm_encode(
            llm,
            queries,
            q_instruction,
        )
134

135
    validate_embed_output(q_rep, d_rep)
136
137
138
139
140
141


@pytest.mark.asyncio
async def test_gritlm_api_server_embedding():
    queries, q_instruction, documents, d_instruction = get_test_data()

142
    args = ["--runner", "pooling", "--max_model_len", str(MAX_MODEL_LEN)]
143

144
    with RemoteOpenAIServer(MODEL_NAME, args) as server:
145
        client_embedding = server.get_async_client()
146

147
148
        d_rep = await run_client_embeddings(
            client_embedding,
149
150
151
            documents,
            d_instruction,
        )
152
153
        q_rep = await run_client_embeddings(
            client_embedding,
154
155
156
            queries,
            q_instruction,
        )
157

158
    validate_embed_output(q_rep, d_rep)
159
160


161
def test_gritlm_offline_generate(monkeypatch: pytest.MonkeyPatch, vllm_runner):
162
    input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"
163

164
    with vllm_runner(
165
166
167
        MODEL_NAME,
        runner="generate",
        max_model_len=MAX_MODEL_LEN,
168
    ) as vllm_model:
169
        llm = vllm_model.llm
170

171
172
        sampling_params = SamplingParams(temperature=0.0, max_tokens=256)
        outputs = llm.generate(input, sampling_params=sampling_params)
173

174
    assert outputs[0].outputs[0].text == "The capital of France is Paris."
175
176
177


@pytest.mark.asyncio
178
async def test_gritlm_api_server_generate():
179
180
    input = "<|user|>\nWhat is the capital of France?\n<|assistant|>\n"

181
    args = ["--runner", "generate", "--max_model_len", str(MAX_MODEL_LEN)]
182

183
    with RemoteOpenAIServer(MODEL_NAME, args) as server:
184
185
186
187
188
189
190
191
        client_generate = server.get_async_client()

        outputs = await client_generate.completions.create(
            model=MODEL_NAME,
            prompt=input,
            max_tokens=256,
            temperature=0.0,
        )
192
193

    assert outputs.choices[0].text == "The capital of France is Paris."