Add vLLM v0.18.1 source tree with KV transfer abort fix
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
0
third_party/vllm/tests/samplers/__init__.py
vendored
Normal file
0
third_party/vllm/tests/samplers/__init__.py
vendored
Normal file
225
third_party/vllm/tests/samplers/test_beam_search.py
vendored
Normal file
225
third_party/vllm/tests/samplers/test_beam_search.py
vendored
Normal file
@@ -0,0 +1,225 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Compare the outputs of HF and vLLM when using beam search.
|
||||
|
||||
Run `pytest tests/samplers/test_beam_search.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from transformers import AutoModelForSeq2SeqLM
|
||||
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
# Extra engine kwargs needed for numerically deterministic beam search.
|
||||
# On ROCm, floating-point reductions in attention and GEMM kernels are
|
||||
# non-associative and sensitive to batch geometry, so we:
|
||||
# async_scheduling=False – deterministic batch composition
|
||||
# enforce_eager=True – no CUDA-graph padding changing effective size
|
||||
# enable_prefix_caching=False – avoid prefix-sharing side effects
|
||||
# max_num_seqs=1 – fixed batch size across runs
|
||||
# On other platforms these are not needed and the dict is empty.
|
||||
EXTRA_ENGINE_KWARGS: dict = (
|
||||
dict(
|
||||
async_scheduling=False,
|
||||
enforce_eager=True,
|
||||
enable_prefix_caching=False,
|
||||
max_num_seqs=1,
|
||||
)
|
||||
if current_platform.is_rocm()
|
||||
else dict(async_scheduling=False, max_num_seqs=1)
|
||||
)
|
||||
|
||||
# FIXME(zhuohan): The test can not pass if we:
|
||||
# 1. Increase max_tokens to 256.
|
||||
# 2. Increase beam_width to 8.
|
||||
# 3. Use the model "huggyllama/llama-7b".
|
||||
MAX_TOKENS = [64]
|
||||
BEAM_WIDTHS = [4]
|
||||
MM_BEAM_WIDTHS = [2]
|
||||
MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
|
||||
@pytest.mark.parametrize("beam_width", BEAM_WIDTHS)
|
||||
def test_beam_search_single_input(
|
||||
monkeypatch,
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
beam_width: int,
|
||||
) -> None:
|
||||
if current_platform.is_rocm():
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
|
||||
|
||||
example_prompts = example_prompts[:1]
|
||||
with hf_runner(model, dtype=dtype) as hf_model:
|
||||
hf_outputs = hf_model.generate_beam_search(
|
||||
example_prompts, beam_width, max_tokens
|
||||
)
|
||||
|
||||
with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_beam_search(
|
||||
example_prompts, beam_width, max_tokens
|
||||
)
|
||||
|
||||
for i in range(len(example_prompts)):
|
||||
hf_output_ids, hf_output_texts = hf_outputs[i]
|
||||
vllm_output_ids, vllm_output_texts = vllm_outputs[i]
|
||||
for j, (hf_text, vllm_text) in enumerate(
|
||||
zip(hf_output_texts, vllm_output_texts)
|
||||
):
|
||||
print(f">>>{j}-th hf output:")
|
||||
print(hf_text)
|
||||
print(f">>>{j}-th vllm output:")
|
||||
print(vllm_text)
|
||||
assert len(hf_output_ids) == len(vllm_output_ids)
|
||||
for j in range(len(hf_output_ids)):
|
||||
assert hf_output_ids[j] == vllm_output_ids[j], (
|
||||
f"Test{i} output{j}:\nHF: {hf_output_ids}\nvLLM: {vllm_output_ids}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
|
||||
@pytest.mark.parametrize("beam_width", BEAM_WIDTHS)
|
||||
def test_beam_search_with_concurrency_limit(
|
||||
monkeypatch,
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
beam_width: int,
|
||||
) -> None:
|
||||
if current_platform.is_rocm():
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
|
||||
|
||||
# example_prompts[1]&[3]&[7] fails due to unknown reason even without
|
||||
# concurrency limit. skip them for now.
|
||||
example_prompts = example_prompts[:8]
|
||||
concurrency_limit = 2
|
||||
assert len(example_prompts) > concurrency_limit
|
||||
with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
|
||||
outputs_with_limit = vllm_model.generate_beam_search(
|
||||
example_prompts,
|
||||
beam_width,
|
||||
max_tokens,
|
||||
concurrency_limit=concurrency_limit,
|
||||
)
|
||||
outputs_without_limit = []
|
||||
|
||||
for i in range(0, len(example_prompts), concurrency_limit):
|
||||
outputs_without_limit.extend(
|
||||
vllm_model.generate_beam_search(
|
||||
example_prompts[i : i + concurrency_limit],
|
||||
beam_width,
|
||||
max_tokens,
|
||||
)
|
||||
)
|
||||
|
||||
correct = True
|
||||
for i in range(len(example_prompts)):
|
||||
output_ids_with_limit, output_texts_with_limit = outputs_with_limit[i]
|
||||
output_ids_without_limit, output_texts_without_limit = outputs_without_limit[i]
|
||||
for j, (text_with_limit, text_without_limit) in enumerate(
|
||||
zip(output_texts_with_limit, output_texts_without_limit)
|
||||
):
|
||||
print(f">>>{j}-th with limit output:")
|
||||
print(text_with_limit)
|
||||
print(f">>>{j}-th without limit output:")
|
||||
print(text_without_limit)
|
||||
assert len(output_ids_with_limit) == len(output_ids_without_limit)
|
||||
for j in range(len(output_ids_with_limit)):
|
||||
if output_ids_with_limit[j] != output_ids_without_limit[j]:
|
||||
print(
|
||||
f"Test{i} output{j}:\n+limit: {output_ids_with_limit}\n"
|
||||
f"-limit: {output_ids_without_limit}"
|
||||
)
|
||||
correct = False
|
||||
assert correct
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
|
||||
@pytest.mark.parametrize("beam_width", MM_BEAM_WIDTHS)
|
||||
def test_beam_search_passes_multimodal_data(
|
||||
monkeypatch,
|
||||
hf_runner,
|
||||
vllm_runner,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
beam_width: int,
|
||||
) -> None:
|
||||
"""Ensure that beam search passes multimodal data through correctly."""
|
||||
if current_platform.is_rocm():
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
|
||||
|
||||
# NOTE - this test is primarily to check that mm data is passed to beams
|
||||
# correctly. As such, we just need to check one extra modality to make
|
||||
# sure things pass through properly.
|
||||
audios = [AudioAsset("mary_had_lamb").audio_and_sample_rate]
|
||||
model = "Qwen/Qwen2-Audio-7B-Instruct"
|
||||
audio_seq = "<|audio_bos|><|AUDIO|><|audio_eos|>"
|
||||
prompts = [
|
||||
f"<|im_start|>user\n{audio_seq}Can you transcribe this?<|im_end|>\n<|im_start|>assistant\n" # noqa: E501
|
||||
]
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSeq2SeqLM) as hf_model:
|
||||
audio_token_id = hf_model.config.audio_token_index
|
||||
eos_token_id = hf_model.tokenizer.eos_token_id # <|im_end|>
|
||||
hf_outputs = hf_model.generate_beam_search(
|
||||
prompts,
|
||||
beam_width=beam_width,
|
||||
max_tokens=max_tokens,
|
||||
audios=audios,
|
||||
)
|
||||
|
||||
with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_beam_search(
|
||||
prompts,
|
||||
beam_width=beam_width,
|
||||
max_tokens=max_tokens,
|
||||
audios=audios,
|
||||
)
|
||||
|
||||
seq_with_no_audio_toks = lambda seq: [tok for tok in seq if tok != audio_token_id]
|
||||
|
||||
for i in range(len(prompts)):
|
||||
hf_output_ids, hf_output_texts = hf_outputs[i]
|
||||
vllm_output_ids, vllm_output_texts = vllm_outputs[i]
|
||||
|
||||
for j, (hf_text, vllm_text) in enumerate(
|
||||
zip(hf_output_texts, vllm_output_texts)
|
||||
):
|
||||
print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:")
|
||||
print(hf_text)
|
||||
print(f">>>{j}-th vllm output:")
|
||||
print(vllm_text)
|
||||
assert len(hf_output_ids) == len(vllm_output_ids)
|
||||
|
||||
for j in range(len(hf_output_ids)):
|
||||
# Compare everything except for the audio tokens; we do this since
|
||||
# the IDs returned from the transformers helper expands the audio
|
||||
# token to match features, while the vLLM helper maintains the
|
||||
# single audio token in the input text
|
||||
filtered_hf_output_ids = seq_with_no_audio_toks(hf_output_ids[j])
|
||||
filtered_vllm_output_ids = seq_with_no_audio_toks(vllm_output_ids[j])
|
||||
|
||||
# HF output IDs may contain the end of sequence
|
||||
if len(filtered_hf_output_ids) == len(filtered_vllm_output_ids) + 1:
|
||||
assert filtered_hf_output_ids[-1] == eos_token_id
|
||||
filtered_hf_output_ids = filtered_hf_output_ids[:-1]
|
||||
|
||||
assert filtered_hf_output_ids == filtered_vllm_output_ids
|
||||
|
||||
|
||||
# NOTE: encoder/decoder tests are currently located under
|
||||
# tests/models/multimodal/generation/test_whisper.py
|
||||
35
third_party/vllm/tests/samplers/test_ignore_eos.py
vendored
Normal file
35
third_party/vllm/tests/samplers/test_ignore_eos.py
vendored
Normal file
@@ -0,0 +1,35 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Make sure ignore_eos works.
|
||||
|
||||
Run `pytest tests/samplers/test_ignore_eos.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import SamplingParams
|
||||
|
||||
# We also test with llama because it has generation_config to specify EOS
|
||||
# (past regression).
|
||||
MODELS = ["distilbert/distilgpt2", "meta-llama/Llama-3.2-1B"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", [512])
|
||||
def test_ignore_eos(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
dtype: str,
|
||||
max_tokens: int,
|
||||
) -> None:
|
||||
with vllm_runner(model, dtype=dtype) as vllm_model:
|
||||
sampling_params = SamplingParams(max_tokens=max_tokens, ignore_eos=True)
|
||||
|
||||
for prompt in example_prompts:
|
||||
ignore_eos_output = vllm_model.llm.generate(
|
||||
prompt, sampling_params=sampling_params
|
||||
)
|
||||
output_length = len(ignore_eos_output[0].outputs[0].token_ids)
|
||||
assert output_length == max_tokens
|
||||
91
third_party/vllm/tests/samplers/test_logprobs.py
vendored
Normal file
91
third_party/vllm/tests/samplers/test_logprobs.py
vendored
Normal file
@@ -0,0 +1,91 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.logprobs import FlatLogprobs
|
||||
|
||||
MODELS = ["distilbert/distilgpt2"]
|
||||
MAX_TOKENS = 5
|
||||
NUM_TOP_LOGPROBS = 5
|
||||
NUM_PROMPT_LOGPROBS = 7
|
||||
MAX_LOGPROBS = max(NUM_TOP_LOGPROBS, NUM_PROMPT_LOGPROBS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("greedy", [True, False])
|
||||
@pytest.mark.parametrize("flat_logprobs", [True, False])
|
||||
def test_ranks(
|
||||
vllm_runner,
|
||||
model,
|
||||
dtype,
|
||||
greedy,
|
||||
flat_logprobs,
|
||||
example_prompts,
|
||||
):
|
||||
with vllm_runner(model, dtype=dtype, max_logprobs=MAX_LOGPROBS) as vllm_model:
|
||||
tokenizer = vllm_model.llm.get_tokenizer()
|
||||
example_prompt_tokens = [tokenizer.encode(prompt) for prompt in example_prompts]
|
||||
sampling_params = SamplingParams(
|
||||
temperature=0.0 if greedy else 1.0,
|
||||
top_p=1.0,
|
||||
max_tokens=MAX_TOKENS,
|
||||
logprobs=NUM_TOP_LOGPROBS,
|
||||
prompt_logprobs=NUM_PROMPT_LOGPROBS,
|
||||
flat_logprobs=flat_logprobs,
|
||||
)
|
||||
results = vllm_model.generate_w_logprobs(example_prompts, sampling_params)
|
||||
|
||||
assert len(results) == len(example_prompt_tokens)
|
||||
for i, (result, prompt_tokens) in enumerate(zip(results, example_prompt_tokens)):
|
||||
decode_tokens, _, decode_logprobs, prompt_logprobs = result
|
||||
|
||||
# Ensure the return type of logprobs is accurate
|
||||
assert isinstance(prompt_logprobs, FlatLogprobs if flat_logprobs else list)
|
||||
assert isinstance(decode_logprobs, FlatLogprobs if flat_logprobs else list)
|
||||
|
||||
########################
|
||||
# Check prompt logprobs
|
||||
########################
|
||||
assert len(prompt_tokens) == len(prompt_logprobs)
|
||||
# No logprob for first prompt token
|
||||
assert not prompt_logprobs[0]
|
||||
for position, (token, logprobs) in enumerate(
|
||||
zip(prompt_tokens[1:], prompt_logprobs[1:]), start=1
|
||||
):
|
||||
# Ensure logprobs of prompt token is always returned
|
||||
logprob = logprobs.get(token)
|
||||
assert logprob is not None
|
||||
assert logprob.rank >= 1
|
||||
# Ensure # of returned logprobs should be
|
||||
# either NUM_PROMPT_LOGPROBS or NUM_PROMPT_LOGPROBS+1
|
||||
assert NUM_PROMPT_LOGPROBS <= len(logprobs) <= NUM_PROMPT_LOGPROBS + 1
|
||||
# Ensure top NUM_PROMPT_LOGPROBS is always extracted
|
||||
assert set(range(1, NUM_PROMPT_LOGPROBS + 1)).issubset(
|
||||
{logprob.rank for logprob in logprobs.values()}
|
||||
)
|
||||
|
||||
########################
|
||||
# Check sample logprobs
|
||||
########################
|
||||
assert len(decode_tokens) == len(decode_logprobs)
|
||||
for position, (token, logprobs) in enumerate(
|
||||
zip(decode_tokens, decode_logprobs)
|
||||
):
|
||||
# Ensure logprobs of chosen token is always returned
|
||||
logprob = logprobs.get(token)
|
||||
assert logprob is not None
|
||||
if greedy:
|
||||
# For greedy sampling, all chosen logprob should be top ranked
|
||||
assert logprob.rank == 1
|
||||
else:
|
||||
assert logprob.rank >= 1
|
||||
# Ensure # of returned logprobs should be
|
||||
# either NUM_TOP_LOGPROBS or NUM_TOP_LOGPROBS+1
|
||||
assert NUM_TOP_LOGPROBS <= len(logprobs) <= NUM_TOP_LOGPROBS + 1
|
||||
# Ensure top NUM_TOP_LOGPROBS logprobs is always extracted
|
||||
assert set(range(1, NUM_TOP_LOGPROBS + 1)).issubset(
|
||||
{logprob.rank for logprob in logprobs.values()}
|
||||
)
|
||||
185
third_party/vllm/tests/samplers/test_no_bad_words.py
vendored
Normal file
185
third_party/vllm/tests/samplers/test_no_bad_words.py
vendored
Normal file
@@ -0,0 +1,185 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Make sure bad_words works.
|
||||
|
||||
Run `pytest tests/samplers/test_no_bad_words.py`.
|
||||
|
||||
"""
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
|
||||
def _generate(
|
||||
llm: LLM,
|
||||
prompt: str,
|
||||
num_prompt_tokens: int,
|
||||
temperature: float = 0,
|
||||
bad_words: list[str] | None = None,
|
||||
) -> list[int]:
|
||||
sampling_params = SamplingParams(
|
||||
temperature=temperature,
|
||||
bad_words=bad_words,
|
||||
)
|
||||
|
||||
# [([output_token_ids, ], [output_text, ]), ]
|
||||
output = llm.generate([prompt], sampling_params=sampling_params)
|
||||
|
||||
output_token_ids = output[0][0][0][num_prompt_tokens:]
|
||||
# [0] first (and only) request output
|
||||
# [0] token_ids (not text)
|
||||
# [0] first (and only) output completion
|
||||
|
||||
return output_token_ids
|
||||
|
||||
|
||||
class TestOneTokenBadWord:
|
||||
MODEL = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
|
||||
PROMPT = "How old are "
|
||||
TARGET_TOKEN = "mn"
|
||||
|
||||
def setup_method(self, method):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL)
|
||||
|
||||
self.num_prompt_tokens = len(self._encode(self.PROMPT))
|
||||
self.target_token_id = self._encode(
|
||||
self.TARGET_TOKEN, add_special_tokens=False
|
||||
)[0]
|
||||
|
||||
def test_one_token_bad_word(self, vllm_runner):
|
||||
with vllm_runner(self.MODEL) as llm:
|
||||
output_token_ids = self._generate(llm)
|
||||
assert output_token_ids[0] == self.target_token_id
|
||||
|
||||
output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN])
|
||||
assert self.target_token_id not in output_token_ids
|
||||
|
||||
def _generate(self, llm: LLM, bad_words: list[str] | None = None) -> list[int]:
|
||||
return _generate(
|
||||
llm=llm,
|
||||
prompt=self.PROMPT,
|
||||
num_prompt_tokens=self.num_prompt_tokens,
|
||||
bad_words=bad_words,
|
||||
)
|
||||
|
||||
def _encode(self, prompt: str, add_special_tokens: bool = True) -> list[int]:
|
||||
return self.tokenizer(prompt, add_special_tokens=add_special_tokens).input_ids
|
||||
|
||||
|
||||
class TestTwoTokenBadWord:
|
||||
# Another model (with a different tokenizer behaviour)
|
||||
MODEL = "distilbert/distilgpt2"
|
||||
|
||||
PROMPT = "How old are you? I am 10"
|
||||
TARGET_TOKEN1 = "years"
|
||||
TARGET_TOKEN2 = "old"
|
||||
NEIGHBOUR_TOKEN2 = "older"
|
||||
|
||||
def setup_method(self, method):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.MODEL, add_prefix_space=True
|
||||
)
|
||||
|
||||
self.num_prompt_tokens = len(self._encode(self.PROMPT))
|
||||
self.target_token_id1 = self._encode(
|
||||
self.TARGET_TOKEN1, add_special_tokens=False
|
||||
)[0]
|
||||
self.target_token_id2 = self._encode(
|
||||
self.TARGET_TOKEN2, add_special_tokens=False
|
||||
)[0]
|
||||
self.neighbour_token_id2 = self._encode(
|
||||
self.NEIGHBOUR_TOKEN2, add_special_tokens=False
|
||||
)[0]
|
||||
|
||||
def test_two_token_bad_word(self, vllm_runner):
|
||||
with vllm_runner(self.MODEL, dtype="half") as llm:
|
||||
output_token_ids = self._generate(llm)
|
||||
assert output_token_ids[:2] == [
|
||||
self.target_token_id1,
|
||||
self.target_token_id2,
|
||||
]
|
||||
|
||||
output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN1])
|
||||
assert self.target_token_id1 not in output_token_ids
|
||||
|
||||
output_token_ids = self._generate(llm, bad_words=[self.TARGET_TOKEN2])
|
||||
assert output_token_ids[0] == self.target_token_id1
|
||||
assert self.target_token_id2 not in output_token_ids
|
||||
|
||||
output_token_ids = self._generate(
|
||||
llm, bad_words=[f"{self.TARGET_TOKEN1} {self.TARGET_TOKEN2}"]
|
||||
)
|
||||
assert output_token_ids[0] == self.target_token_id1
|
||||
assert output_token_ids[:2] != [
|
||||
self.target_token_id1,
|
||||
self.target_token_id2,
|
||||
]
|
||||
assert not self._contains(
|
||||
output_token_ids, [self.target_token_id1, self.target_token_id2]
|
||||
)
|
||||
# Model dependent behaviour
|
||||
assert output_token_ids[:2] == [
|
||||
self.target_token_id1,
|
||||
self.neighbour_token_id2,
|
||||
]
|
||||
|
||||
output_token_ids = self._generate(
|
||||
llm,
|
||||
bad_words=[
|
||||
f"{self.TARGET_TOKEN1} {self.TARGET_TOKEN2}",
|
||||
f"{self.TARGET_TOKEN1} {self.NEIGHBOUR_TOKEN2}",
|
||||
],
|
||||
)
|
||||
assert output_token_ids[0] == self.target_token_id1
|
||||
assert output_token_ids[:2] != [
|
||||
self.target_token_id1,
|
||||
self.target_token_id2,
|
||||
]
|
||||
assert not self._contains(
|
||||
output_token_ids, [self.target_token_id1, self.target_token_id2]
|
||||
)
|
||||
assert output_token_ids[:2] != [
|
||||
self.target_token_id1,
|
||||
self.neighbour_token_id2,
|
||||
]
|
||||
assert not self._contains(
|
||||
output_token_ids, [self.target_token_id1, self.neighbour_token_id2]
|
||||
)
|
||||
assert (self.target_token_id2 in output_token_ids) or (
|
||||
self.neighbour_token_id2 in output_token_ids
|
||||
)
|
||||
|
||||
def _generate(self, llm: LLM, bad_words: list[str] | None = None) -> list[int]:
|
||||
return _generate(
|
||||
llm=llm,
|
||||
prompt=self.PROMPT,
|
||||
num_prompt_tokens=self.num_prompt_tokens,
|
||||
bad_words=bad_words,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _contains(sequence: list[int], subsequence: list[int]) -> bool:
|
||||
searched = False
|
||||
|
||||
for start in range(len(sequence)):
|
||||
end = start + len(subsequence)
|
||||
current_subsequence = sequence[start:end]
|
||||
|
||||
if len(current_subsequence) < len(subsequence):
|
||||
continue
|
||||
|
||||
searched = True
|
||||
|
||||
assert len(current_subsequence) == len(subsequence)
|
||||
|
||||
if current_subsequence == subsequence:
|
||||
return True
|
||||
|
||||
assert searched, "All subsequences did not match in length..."
|
||||
|
||||
return False
|
||||
|
||||
def _encode(self, prompt: str, add_special_tokens: bool = True) -> list[int]:
|
||||
return self.tokenizer(prompt, add_special_tokens=add_special_tokens).input_ids
|
||||
Reference in New Issue
Block a user