Skip to content

Commit 96df0e7

Browse files
clefourrierclefourrier
committed
update caching"
1 parent 951cbc0 commit 96df0e7

12 files changed

Lines changed: 127 additions & 79 deletions

File tree

‎docs/source/evaluating-a-custom-model.mdx‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ Here's a basic example:
1414
from typing import List
1515
from lighteval.models.abstract_model import LightevalModel
1616
from lighteval.models.model_output import ModelResponse
17-
from lighteval.tasks.requests import Doc
17+
from lighteval.tasks.requests import Doc, SamplingMethod
1818
from lighteval.utils.cache_management import SampleCache, cached
1919

2020
class MyCustomModel(LightevalModel):
@@ -25,17 +25,17 @@ class MyCustomModel(LightevalModel):
2525
# Enable caching (recommended)
2626
self._cache = SampleCache(config)
2727

28-
@cached("predictions") # Enable caching for better performance
28+
@cached("predictions", SamplingMethod.GENERATIVE) # Enable caching for better performance
2929
def greedy_until(self, docs: List[Doc]) -> List[ModelResponse]:
3030
# Implement generation logic
3131
pass
3232

33-
@cached("predictions") # Enable caching for better performance
33+
@cached("loglikelihood", SamplingMethod.LOGPROBS) # Enable caching for better performance
3434
def loglikelihood(self, docs: List[Doc]) -> List[ModelResponse]:
3535
# Implement loglikelihood computation
3636
pass
3737

38-
@cached("predictions") # Enable caching for better performance
38+
@cached("loglikelihood", SamplingMethod.LOGPROBS) # Enable caching for better performance
3939
def loglikelihood_rolling(self, docs: List[Doc]) -> List[ModelResponse]:
4040
# Implement rolling loglikelihood computation
4141
pass
@@ -130,7 +130,7 @@ To enable caching in your custom model:
130130

131131
3. **Add cache decorators** to your prediction methods:
132132
```python
133-
@cached("predictions")
133+
@cached("predictions", SamplingMethod.GENERATIVE)
134134
def greedy_until(self, docs: List[Doc]) -> List[ModelResponse]:
135135
# Your implementation...
136136
```

‎src/lighteval/models/dummy/dummy_model.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828

2929
from lighteval.models.abstract_model import LightevalModel, ModelConfig
3030
from lighteval.models.model_output import ModelResponse
31-
from lighteval.tasks.requests import Doc
31+
from lighteval.tasks.requests import Doc, SamplingMethod
3232
from lighteval.utils.cache_management import SampleCache, cached
3333

3434

@@ -88,11 +88,11 @@ def add_special_tokens(self):
8888
def max_length(self) -> int:
8989
return 2048
9090

91-
@cached("predictions")
91+
@cached("predictions", SamplingMethod.GENERATIVE)
9292
def greedy_until(self, docs: list[Doc]) -> list[ModelResponse]:
9393
return [ModelResponse(text=["random baseline"]) for _ in range(len(docs))]
9494

95-
@cached("predictions")
95+
@cached("predictions", SamplingMethod.LOGPROBS)
9696
def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
9797
model_responses = []
9898
for doc in docs:
@@ -105,7 +105,7 @@ def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
105105

106106
return model_responses
107107

108-
@cached("predictions")
108+
@cached("predictions", SamplingMethod.LOGPROBS)
109109
def loglikelihood_rolling(self, docs: list[Doc]) -> list[ModelResponse]:
110110
model_responses = []
111111
for doc in docs:

‎src/lighteval/models/endpoints/endpoint_model.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@
4848
from lighteval.models.abstract_model import LightevalModel, ModelConfig
4949
from lighteval.models.model_output import ModelResponse
5050
from lighteval.tasks.prompt_manager import PromptManager
51-
from lighteval.tasks.requests import Doc
51+
from lighteval.tasks.requests import Doc, SamplingMethod
5252
from lighteval.utils.cache_management import SampleCache, cached
5353

5454

@@ -545,7 +545,7 @@ def _process_batch_logprob(self, docs: list[Doc], rolling: bool = False) -> list
545545
for context, doc in zip(contexts, docs)
546546
]
547547

548-
@cached("predictions")
548+
@cached("predictions", SamplingMethod.GENERATIVE)
549549
def greedy_until(
550550
self,
551551
docs: List[Doc],
@@ -589,11 +589,11 @@ def _greedy_until(self, docs: List[Doc]) -> list[ModelResponse]:
589589

590590
return dataset.get_original_order(results)
591591

592-
@cached("predictions")
592+
@cached("predictions", SamplingMethod.LOGPROBS)
593593
def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
594594
return self._loglikelihood(docs, rolling=False)
595595

596-
@cached("predictions")
596+
@cached("predictions", SamplingMethod.LOGPROBS)
597597
def loglikelihood_rolling(self, docs: list[Doc], override_bs=None) -> list[ModelResponse]:
598598
return self._loglikelihood(docs, rolling=True)
599599

‎src/lighteval/models/endpoints/inference_providers_model.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
from lighteval.models.abstract_model import LightevalModel, ModelConfig
3636
from lighteval.models.model_output import ModelResponse
3737
from lighteval.tasks.prompt_manager import PromptManager
38-
from lighteval.tasks.requests import Doc
38+
from lighteval.tasks.requests import Doc, SamplingMethod
3939
from lighteval.utils.cache_management import SampleCache, cached
4040

4141

@@ -191,7 +191,7 @@ async def bounded_api_call(prompt, num_samples):
191191

192192
return results
193193

194-
@cached("predictions")
194+
@cached("predictions", SamplingMethod.GENERATIVE)
195195
def greedy_until(
196196
self,
197197
docs: list[Doc],
@@ -250,14 +250,14 @@ def max_length(self) -> int:
250250
logger.warning("Tokenizer was not correctly loaded. Max model context length is assumed to be 30K tokens")
251251
return 30000
252252

253-
@cached("predictions")
253+
@cached("predictions", SamplingMethod.LOGPROBS)
254254
def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
255255
"""Tokenize the context and continuation and compute the log likelihood of those
256256
tokenized sequences.
257257
"""
258258
raise NotImplementedError
259259

260-
@cached("predictions")
260+
@cached("predictions", SamplingMethod.LOGPROBS)
261261
def loglikelihood_rolling(self, docs: list[Doc]) -> list[ModelResponse]:
262262
"""This function is used to compute the log likelihood of the context for perplexity metrics."""
263263
raise NotImplementedError

‎src/lighteval/models/endpoints/litellm_model.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
from lighteval.models.abstract_model import LightevalModel, ModelConfig
3131
from lighteval.models.model_output import ModelResponse
3232
from lighteval.tasks.prompt_manager import PromptManager
33-
from lighteval.tasks.requests import Doc
33+
from lighteval.tasks.requests import Doc, SamplingMethod
3434
from lighteval.utils.cache_management import SampleCache, cached
3535
from lighteval.utils.imports import is_litellm_available
3636

@@ -254,7 +254,7 @@ def __call_api_parallel(
254254

255255
return results
256256

257-
@cached("predictions")
257+
@cached("predictions", SamplingMethod.GENERATIVE)
258258
def greedy_until(
259259
self,
260260
docs: list[Doc],
@@ -321,14 +321,14 @@ def max_length(self) -> int:
321321
"""Return the maximum sequence length of the model."""
322322
return 4096
323323

324-
@cached("predictions")
324+
@cached("predictions", SamplingMethod.LOGPROBS)
325325
def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
326326
"""Tokenize the context and continuation and compute the log likelihood of those
327327
tokenized sequences.
328328
"""
329329
raise NotImplementedError
330330

331-
@cached("predictions")
331+
@cached("predictions", SamplingMethod.LOGPROBS)
332332
def loglikelihood_rolling(self, docs: list[Doc]) -> list[ModelResponse]:
333333
"""This function is used to compute the log likelihood of the context for perplexity metrics."""
334334
raise NotImplementedError

‎src/lighteval/models/nanotron/nanotron_model.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
from lighteval.models.transformers.transformers_model import LightevalModel
4949
from lighteval.tasks.requests import (
5050
Doc,
51+
SamplingMethod,
5152
)
5253
from lighteval.utils.cache_management import SampleCache, cached
5354
from lighteval.utils.imports import is_nanotron_available
@@ -473,7 +474,7 @@ def _check_continuations_start_space(self, continuation: str) -> str:
473474
continuation = continuation.lstrip()
474475
return continuation
475476

476-
@cached("predictions")
477+
@cached("predictions", SamplingMethod.LOGPROBS)
477478
def loglikelihood(self, requests: List[Doc]) -> List[ModelResponse]:
478479
"""Tokenize the context and continuation and compute the log likelihood of those
479480
tokenized sequences.
@@ -496,7 +497,7 @@ def loglikelihood(self, requests: List[Doc]) -> List[ModelResponse]:
496497
disable_tqdm=bool(dist.get_rank(self.parallel_context.world_pg) != 0),
497498
)
498499

499-
@cached("predictions")
500+
@cached("predictions", SamplingMethod.LOGPROBS)
500501
def loglikelihood_rolling(self, requests: List[Doc]) -> List[ModelResponse]:
501502
"""This function is used to compute the log likelihood of the context for perplexity metrics."""
502503
for request in tqdm(
@@ -931,7 +932,7 @@ def _loglikelihood_tokens(
931932
return dataset.get_original_order(res)
932933

933934
@torch.inference_mode()
934-
@cached("predictions")
935+
@cached("predictions", SamplingMethod.GENERATIVE)
935936
def greedy_until(
936937
self,
937938
requests: List[Doc],

‎src/lighteval/models/sglang/sglang_model.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
from lighteval.models.model_output import ModelResponse
3434
from lighteval.models.utils import _simplify_name, uses_chat_template
3535
from lighteval.tasks.prompt_manager import PromptManager
36-
from lighteval.tasks.requests import Doc
36+
from lighteval.tasks.requests import Doc, SamplingMethod
3737
from lighteval.utils.cache_management import SampleCache, cached
3838
from lighteval.utils.imports import is_sglang_available
3939

@@ -216,7 +216,7 @@ def _create_auto_tokenizer(self, config: SGLangModelConfig):
216216
tokenizer.pad_token = tokenizer.eos_token
217217
return tokenizer
218218

219-
@cached("predictions")
219+
@cached("predictions", SamplingMethod.GENERATIVE)
220220
def greedy_until(
221221
self,
222222
docs: list[Doc],
@@ -345,7 +345,7 @@ def _generate(
345345
)
346346
return outputs
347347

348-
@cached("predictions")
348+
@cached("predictions", SamplingMethod.LOGPROBS)
349349
def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
350350
return self._loglikelihood_tokens(docs)
351351

@@ -414,6 +414,6 @@ def _loglikelihood_tokens(
414414
res.append(answer)
415415
return dataset.get_original_order(res)
416416

417-
@cached("predictions")
417+
@cached("predictions", SamplingMethod.LOGPROBS)
418418
def loglikelihood_rolling(self, docs: list[Doc]) -> list[ModelResponse]:
419419
raise NotImplementedError()

‎src/lighteval/models/transformers/transformers_model.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@
5252
)
5353
from lighteval.models.utils import _get_dtype, _get_model_sha, _simplify_name, uses_chat_template
5454
from lighteval.tasks.prompt_manager import PromptManager
55-
from lighteval.tasks.requests import Doc
55+
from lighteval.tasks.requests import Doc, SamplingMethod
5656
from lighteval.utils.cache_management import SampleCache, cached
5757
from lighteval.utils.imports import (
5858
is_accelerate_available,
@@ -740,7 +740,7 @@ def _padded_greedy_until(
740740

741741
return dataset.get_original_order(results)
742742

743-
@cached("predictions")
743+
@cached("predictions", SamplingMethod.GENERATIVE)
744744
def greedy_until(
745745
self,
746746
docs: list[Doc],
@@ -867,7 +867,7 @@ def _generate(
867867
else:
868868
return self._generate_padded(**kwargs)
869869

870-
@cached("predictions")
870+
@cached("predictions", SamplingMethod.LOGPROBS)
871871
def loglikelihood(
872872
self,
873873
docs: list[Doc],
@@ -883,7 +883,7 @@ def loglikelihood(
883883
"""
884884
return self._loglikelihood_tokens(docs)
885885

886-
@cached("predictions")
886+
@cached("predictions", SamplingMethod.LOGPROBS)
887887
def loglikelihood_rolling(
888888
self,
889889
docs: list[Doc],

‎src/lighteval/models/transformers/vlm_transformers_model.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
from lighteval.models.model_output import ModelResponse
4545
from lighteval.models.utils import _get_dtype, _get_model_sha, _simplify_name
4646
from lighteval.tasks.prompt_manager import PromptManager
47-
from lighteval.tasks.requests import Doc
47+
from lighteval.tasks.requests import Doc, SamplingMethod
4848
from lighteval.utils.cache_management import SampleCache, cached
4949
from lighteval.utils.imports import (
5050
is_accelerate_available,
@@ -333,7 +333,7 @@ def _init_max_length(self) -> int:
333333

334334
return 2048
335335

336-
@cached("predictions")
336+
@cached("predictions", SamplingMethod.GENERATIVE)
337337
def greedy_until(
338338
self,
339339
docs: list[Doc],
@@ -424,14 +424,14 @@ def _greedy_until(
424424

425425
return dataset.get_original_order(results)
426426

427-
@cached("predictions")
427+
@cached("predictions", SamplingMethod.LOGPROBS)
428428
def loglikelihood(
429429
self,
430430
docs: list[Doc],
431431
) -> list[ModelResponse]:
432432
raise NotImplementedError()
433433

434-
@cached("predictions")
434+
@cached("predictions", SamplingMethod.LOGPROBS)
435435
def loglikelihood_rolling(
436436
self,
437437
docs: list[Doc],

‎src/lighteval/models/vllm/vllm_model.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@
3636
from lighteval.models.model_output import ModelResponse
3737
from lighteval.models.utils import _simplify_name, uses_chat_template
3838
from lighteval.tasks.prompt_manager import PromptManager
39-
from lighteval.tasks.requests import Doc
39+
from lighteval.tasks.requests import Doc, SamplingMethod
4040
from lighteval.utils.cache_management import SampleCache, cached
4141
from lighteval.utils.imports import is_vllm_available
4242

@@ -259,6 +259,7 @@ def _create_auto_model(self, config: VLLMModelConfig) -> Optional[LLM]:
259259
"seed": int(config.seed),
260260
"max_num_seqs": int(config.max_num_seqs),
261261
"max_num_batched_tokens": int(config.max_num_batched_tokens),
262+
"enforce_eager": True,
262263
}
263264

264265
if config.quantization is not None:
@@ -300,7 +301,7 @@ def _create_auto_tokenizer(self, config: VLLMModelConfig):
300301
tokenizer.pad_token = tokenizer.eos_token
301302
return tokenizer
302303

303-
@cached("predictions")
304+
@cached("predictions", SamplingMethod.GENERATIVE)
304305
def greedy_until(
305306
self,
306307
docs: list[Doc],
@@ -459,7 +460,7 @@ def run_inference_one_model(model_args: dict, sampling_params: SamplingParams, r
459460

460461
return outputs
461462

462-
@cached("predictions")
463+
@cached("predictions", SamplingMethod.LOGPROBS)
463464
def loglikelihood(self, docs: list[Doc]) -> list[ModelResponse]:
464465
return self._loglikelihood_tokens(docs)
465466

@@ -528,7 +529,7 @@ def _loglikelihood_tokens(
528529

529530
return dataset.get_original_order(res)
530531

531-
@cached("predictions")
532+
@cached("predictions", SamplingMethod.LOGPROBS)
532533
def loglikelihood_rolling(self, docs: list[Doc]) -> list[ModelResponse]:
533534
raise NotImplementedError()
534535

@@ -624,7 +625,7 @@ async def _async_batch(self, docs: list[Doc], generative: bool) -> list:
624625
results = await asyncio.gather(*processed_requests)
625626
return results
626627

627-
@cached("predictions")
628+
@cached("predictions", SamplingMethod.GENERATIVE)
628629
async def greedy_until(
629630
self,
630631
docs: list[Doc],
@@ -659,7 +660,7 @@ async def greedy_until(
659660

660661
return results
661662

662-
@cached("predictions")
663+
@cached("predictions", SamplingMethod.LOGPROBS)
663664
async def loglikelihood(
664665
self,
665666
docs: list[Doc],

0 commit comments

Comments
 (0)