zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-12-16 02:08:55 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Huge speed improvement in validation mode.

Browse Source
This commit is contained in:
Goekdeniz-Guelmez
2025-02-21 22:08:49 +01:00
parent 2f20107d9b
commit 6086137131
2 changed files with 35 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
llms/mlx_lm/tuner/datasets.py
View File

@@ -30,7 +30,7 @@ class GRPODataset:
prompt_tokens = tokenizer.apply_chat_template(
[
{'role': 'system', 'content': """A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think><answer> answer here </answer>."""},
{'role': 'user', 'content': prompt_str}
{'role': 'user', 'content': prompt_str}
],
add_generation_prompt=True
)
@@ -44,11 +44,9 @@ class GRPODataset:
self._data.append((prompt_tokens, answer_tokens, prompt_str, answer_str))
def __getitem__(self, idx: int) -> Tuple[List[int], List[int], str, str]:
"""Returns a (prompt_tokens, answer_tokens, prompt_str, answer_str) tuple."""
return self._data[idx]
def __len__(self) -> int:
"""Returns the number of examples in the dataset."""
return len(self._data)