ai_dataloader/collate/torch_collate/
map.rs

1use super::super::Collate;
2use super::TorchCollate;
3use std::{
4    cmp::Eq,
5    collections::{BTreeMap, HashMap},
6    hash::{BuildHasher, Hash},
7};
8
9impl<K, V, H> Collate<HashMap<K, V, H>> for TorchCollate
10where
11    K: Eq + Hash + Clone,
12    V: Clone,
13    Self: Collate<V>,
14    H: BuildHasher,
15{
16    type Output = HashMap<K, <Self as Collate<V>>::Output>;
17    fn collate(&self, batch: Vec<HashMap<K, V, H>>) -> Self::Output {
18        let mut collated = HashMap::with_capacity(batch[0].keys().len());
19        for key in batch[0].keys() {
20            let vec: Vec<_> = batch.iter().map(|hash_map| hash_map[key].clone()).collect();
21            collated.insert(key.clone(), self.collate(vec));
22        }
23        collated
24    }
25}
26impl<K, V> Collate<BTreeMap<K, V>> for TorchCollate
27where
28    K: Ord + Clone,
29    V: Clone,
30    Self: Collate<V>,
31{
32    type Output = BTreeMap<K, <Self as Collate<V>>::Output>;
33    fn collate(&self, batch: Vec<BTreeMap<K, V>>) -> Self::Output {
34        let mut collated = BTreeMap::new();
35        for key in batch[0].keys() {
36            let vec: Vec<_> = batch.iter().map(|hash_map| hash_map[key].clone()).collect();
37            collated.insert(key.clone(), self.collate(vec));
38        }
39        collated
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46    use tch::Tensor;
47
48    #[test]
49    fn vec_of_hash_map() {
50        let map1 = HashMap::from([("A", 0), ("B", 1)]);
51        let map2 = HashMap::from([("A", 100), ("B", 100)]);
52        let expected_result = HashMap::from([
53            ("A", Tensor::from_slice(&[0, 100])),
54            ("B", Tensor::from_slice(&[1, 100])),
55        ]);
56        assert_eq!(TorchCollate.collate(vec![map1, map2]), expected_result);
57
58        // Same value type but different key
59        let map1 = HashMap::from([(1, 0), (2, 1)]);
60        let map2 = HashMap::from([(1, 100), (2, 100)]);
61        let expected_result = HashMap::from([
62            (1, Tensor::from_slice(&[0, 100])),
63            (2, Tensor::from_slice(&[1, 100])),
64        ]);
65        assert_eq!(TorchCollate.collate(vec![map1, map2]), expected_result);
66
67        let map1 = HashMap::from([("A", 0.0), ("B", 1.0)]);
68        let map2 = HashMap::from([("A", 100.0), ("B", 100.0)]);
69        let expected_result = HashMap::from([
70            ("A", Tensor::from_slice(&[0.0, 100.0])),
71            ("B", Tensor::from_slice(&[1.0, 100.0])),
72        ]);
73        assert_eq!(TorchCollate.collate(vec![map1, map2]), expected_result);
74    }
75
76    #[test]
77    fn specialized() {
78        let map1 = HashMap::from([("A", String::from("0")), ("B", String::from("1"))]);
79        let map2 = HashMap::from([("A", String::from("100")), ("B", String::from("100"))]);
80        let expected_result = HashMap::from([
81            ("A", vec![String::from("0"), String::from("100")]),
82            ("B", vec![String::from("1"), String::from("100")]),
83        ]);
84        assert_eq!(TorchCollate.collate(vec![map1, map2]), expected_result);
85
86        let map1 = HashMap::from([("A", "0"), ("B", "1")]);
87        let map2 = HashMap::from([("A", "100"), ("B", "100")]);
88        let expected_result = HashMap::from([("A", vec!["0", "100"]), ("B", vec!["1", "100"])]);
89        assert_eq!(TorchCollate.collate(vec![map1, map2]), expected_result);
90    }
91}
ai_dataloader/collate/torch_collate/map.rs

ai_dataloader/collate/torch_collate/
map.rs
