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

Generalize the get_item method to all CompletionDatasets

Browse Source
This commit is contained in:
Chime Ogbuji 2024-12-09 11:46:00 -05:00 committed by Awni Hannun
parent 7989d0a874
commit b9748e9ee4
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -108,6 +108,14 @@ class CompletionsDatasetCollection:
return self.__fetch_and_process_item__(idx, getitem)
def get_item(
self, idx: int, tokenize: bool = False, add_generation_prompt: bool = True
) -> str:
def getitem(dataset: CompletionsDataset, index: int):
return dataset.get_item(index, tokenize, add_generation_prompt)
return self.__fetch_and_process_item__(idx, getitem)
def get_prompt_and_completion(self, idx: int):
def getitem(dataset: CompletionsDataset, index: int):
return dataset.get_prompt_and_completion(index)