burn_dataset/transform/
mapper.rs

1use crate::Dataset;
2use std::error::Error;
3use std::marker::PhantomData;
4
5/// Basic mapper trait to be used with the [mapper dataset](MapperDataset).
6pub trait Mapper<I, O>: Send + Sync {
7    /// Maps an item of type I to an item of type O.
8    fn map(&self, item: &I) -> O;
9}
10
11/// Dataset mapping each element in an inner dataset to another element type lazily.
12#[derive(new)]
13pub struct MapperDataset<D, M, I> {
14    dataset: D,
15    mapper: M,
16    input: PhantomData<I>,
17}
18
19impl<D, M, I, O, E> Dataset<O, E> for MapperDataset<D, M, I>
20where
21    D: Dataset<I, E>,
22    M: Mapper<I, O> + Send + Sync,
23    I: Send + Sync,
24    O: Send + Sync,
25    E: Error + Send + Sync + 'static,
26{
27    fn get(&self, index: usize) -> Result<O, E> {
28        let item = self.dataset.get(index)?;
29        Ok(self.mapper.map(&item))
30    }
31
32    fn len(&self) -> usize {
33        self.dataset.len()
34    }
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40    use crate::{InMemDataset, test_data};
41
42    #[test]
43    pub fn given_mapper_dataset_when_iterate_should_iterate_though_all_map_items() {
44        struct StringToFirstChar;
45
46        impl Mapper<String, String> for StringToFirstChar {
47            fn map(&self, item: &String) -> String {
48                let mut item = item.clone();
49                item.truncate(1);
50                item
51            }
52        }
53
54        let items_original = test_data::string_items();
55        let dataset = InMemDataset::new(items_original);
56        let dataset = MapperDataset::new(dataset, StringToFirstChar);
57
58        let items: Vec<String> = dataset.iter().map(Result::unwrap).collect();
59
60        assert_eq!(vec!["1", "2", "3", "4"], items);
61    }
62}
burn_dataset/transform/mapper.rs

burn_dataset/transform/
mapper.rs
