| 1412 | self.assertEqual(dset[0], {"en": "aa", "fr": "cc", "translation": {"en": "aa", "fr": "cc"}}) |
| 1413 | |
| 1414 | def test_map_fn_kwargs(self, in_memory): |
| 1415 | with tempfile.TemporaryDirectory() as tmp_dir: |
| 1416 | with Dataset.from_dict({"id": range(10)}) as dset: |
| 1417 | with self._to(in_memory, tmp_dir, dset) as dset: |
| 1418 | fn_kwargs = {"offset": 3} |
| 1419 | with dset.map( |
| 1420 | lambda example, offset: {"id+offset": example["id"] + offset}, fn_kwargs=fn_kwargs |
| 1421 | ) as mapped_dset: |
| 1422 | assert mapped_dset["id+offset"] == list(range(3, 13)) |
| 1423 | with dset.map( |
| 1424 | lambda id, offset: {"id+offset": id + offset}, fn_kwargs=fn_kwargs, input_columns="id" |
| 1425 | ) as mapped_dset: |
| 1426 | assert mapped_dset["id+offset"] == list(range(3, 13)) |
| 1427 | with dset.map( |
| 1428 | lambda id, i, offset: {"id+offset": i + offset}, |
| 1429 | fn_kwargs=fn_kwargs, |
| 1430 | input_columns="id", |
| 1431 | with_indices=True, |
| 1432 | ) as mapped_dset: |
| 1433 | assert mapped_dset["id+offset"] == list(range(3, 13)) |
| 1434 | |
| 1435 | def test_map_caching(self, in_memory): |
| 1436 | with tempfile.TemporaryDirectory() as tmp_dir: |