Implement conversion to explicit subword embeddings.

Author: sebpuetz

src/chunks/storage/array.rs

@@ -5,7 +5,7 @@ use std::mem::size_of;
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use ndarray::{Array2, ArrayView2, ArrayViewMut2, Axis, CowArray, Ix1};
 
-use super::{Storage, StorageView, StorageViewMut};
+use super::{sealed::CloneFromMapping, Storage, StorageView, StorageViewMut};
 use crate::chunks::io::{ChunkIdentifier, ReadChunk, TypeId, WriteChunk};
 use crate::error::{Error, Result};
 use crate::util::padding;
@@ -27,14 +27,13 @@ mod mmap {
     use ndarray::{Array2, ArrayView2, Axis, CowArray, Ix1};
     use ndarray::{Dimension, Ix2};
 
-    #[cfg(target_endian = "little")]
     use super::NdArray;
     #[cfg(target_endian = "little")]
     use crate::chunks::io::WriteChunk;
     use crate::chunks::io::{ChunkIdentifier, TypeId};
-    use crate::chunks::storage::Storage;
     #[cfg(target_endian = "little")]
     use crate::chunks::storage::StorageView;
+    use crate::chunks::storage::{sealed::CloneFromMapping, Storage};
     use crate::error::{Error, Result};
     use crate::util::padding;
 
@@ -103,6 +102,14 @@ mod mmap {
         }
     }
 
+    impl CloneFromMapping for MmapArray {
+        type Result = NdArray;
+
+        fn clone_from_mapping(&self, mapping: &[usize]) -> Self::Result {
+            NdArray::new(self.embeddings(mapping))
+        }
+    }
+
     impl MmapChunk for MmapArray {
         fn mmap_chunk(read: &mut BufReader<File>) -> Result<Self> {
             ChunkIdentifier::ensure_chunk_type(read, ChunkIdentifier::NdArray)?;
@@ -287,6 +294,14 @@ impl StorageViewMut for NdArray {
     }
 }
 
+impl CloneFromMapping for NdArray {
+    type Result = NdArray;
+
+    fn clone_from_mapping(&self, mapping: &[usize]) -> Self::Result {
+        NdArray::new(self.embeddings(mapping))
+    }
+}
+
 impl ReadChunk for NdArray {
     fn read_chunk<R>(read: &mut R) -> Result<Self>
     where

src/chunks/storage/mod.rs

@@ -29,6 +29,20 @@ pub trait Storage {
     fn shape(&self) -> (usize, usize);
 }
 
+pub(crate) mod sealed {
+    use crate::chunks::storage::Storage;
+
+    /// Return a new storage from an existing Storage based on a mapping.
+    pub trait CloneFromMapping {
+        type Result: Storage;
+
+        /// Construct a new Storage based on a mapping.
+        ///
+        /// The `i`th entry in the returned storage is based on `self.embedding(mapping[i])`.
+        fn clone_from_mapping(&self, mapping: &[usize]) -> Self::Result;
+    }
+}
+
 /// Storage that provide a view of the embedding matrix.
 pub trait StorageView: Storage {
     /// Get a view of the embedding matrix.

src/chunks/storage/quantized.rs

@@ -8,7 +8,7 @@ use rand::{RngCore, SeedableRng};
 use rand_xorshift::XorShiftRng;
 use reductive::pq::{QuantizeVector, ReconstructVector, TrainPQ, PQ};
 
-use super::{Storage, StorageView};
+use super::{sealed::CloneFromMapping, Storage, StorageView};
 use crate::chunks::io::{ChunkIdentifier, ReadChunk, TypeId, WriteChunk};
 use crate::error::{Error, Result};
 use crate::storage::NdArray;
@@ -280,6 +280,20 @@ impl Storage for QuantizedArray {
     }
 }
 
+impl CloneFromMapping for QuantizedArray {
+    type Result = QuantizedArray;
+
+    fn clone_from_mapping(&self, mapping: &[usize]) -> Self::Result {
+        let quantized_embeddings = self.quantized_embeddings.select(Axis(0), &mapping);
+        let norms = self.norms.as_ref().map(|n| n.select(Axis(0), &mapping));
+        QuantizedArray {
+            quantizer: self.quantizer.clone(),
+            quantized_embeddings,
+            norms,
+        }
+    }
+}
+
 impl ReadChunk for QuantizedArray {
     fn read_chunk<R>(read: &mut R) -> Result<Self>
     where
@@ -472,7 +486,7 @@ mod mmap {
     use super::{PQRead, QuantizedArray, Storage};
     use crate::chunks::io::MmapChunk;
     use crate::chunks::io::{ChunkIdentifier, WriteChunk};
-    use crate::chunks::storage::NdArray;
+    use crate::chunks::storage::{sealed::CloneFromMapping, NdArray};
     use crate::error::{Error, Result};
     use byteorder::{LittleEndian, ReadBytesExt};
 
@@ -570,6 +584,21 @@ mod mmap {
         }
     }
 
+    impl CloneFromMapping for MmapQuantizedArray {
+        type Result = QuantizedArray;
+
+        fn clone_from_mapping(&self, mapping: &[usize]) -> Self::Result {
+            let quantized_embeddings =
+                unsafe { self.quantized_embeddings() }.select(Axis(0), &mapping);
+            let norms = self.norms.as_ref().map(|n| n.select(Axis(0), &mapping));
+            QuantizedArray {
+                quantizer: self.quantizer.clone(),
+                quantized_embeddings,
+                norms,
+            }
+        }
+    }
+
     impl MmapChunk for MmapQuantizedArray {
         fn mmap_chunk(read: &mut BufReader<File>) -> Result<Self> {
             ChunkIdentifier::ensure_chunk_type(read, ChunkIdentifier::QuantizedArray)?;

src/chunks/storage/wrappers.rs

@@ -14,6 +14,7 @@ use super::{NdArray, QuantizedArray, Storage, StorageView};
 #[cfg(feature = "memmap")]
 use crate::chunks::io::MmapChunk;
 use crate::chunks::io::{ChunkIdentifier, ReadChunk, WriteChunk};
+use crate::chunks::storage::sealed::CloneFromMapping;
 use crate::error::{Error, Result};
 
 /// Storage types wrapper.
@@ -72,6 +73,27 @@ impl Storage for StorageWrap {
     }
 }
 
+impl CloneFromMapping for StorageWrap {
+    type Result = StorageWrap;
+
+    fn clone_from_mapping(&self, mapping: &[usize]) -> Self::Result {
+        match self {
+            StorageWrap::QuantizedArray(quant) => {
+                StorageWrap::QuantizedArray(Box::new(quant.clone_from_mapping(mapping)))
+            }
+            #[cfg(feature = "memmap")]
+            StorageWrap::MmapQuantizedArray(quant) => {
+                StorageWrap::QuantizedArray(Box::new(quant.clone_from_mapping(mapping)))
+            }
+            #[cfg(feature = "memmap")]
+            StorageWrap::MmapArray(mmapped) => {
+                StorageWrap::NdArray(mmapped.clone_from_mapping(mapping))
+            }
+            StorageWrap::NdArray(array) => StorageWrap::NdArray(array.clone_from_mapping(mapping)),
+        }
+    }
+}
+
 #[cfg(feature = "memmap")]
 impl From<MmapArray> for StorageWrap {
     fn from(s: MmapArray) -> Self {
@@ -236,6 +258,22 @@ impl StorageView for StorageViewWrap {
     }
 }
 
+impl CloneFromMapping for StorageViewWrap {
+    type Result = StorageViewWrap;
+
+    fn clone_from_mapping(&self, mapping: &[usize]) -> Self::Result {
+        match self {
+            #[cfg(all(feature = "memmap", target_endian = "little"))]
+            StorageViewWrap::MmapArray(mmapped) => {
+                StorageViewWrap::NdArray(mmapped.clone_from_mapping(mapping))
+            }
+            StorageViewWrap::NdArray(array) => {
+                StorageViewWrap::NdArray(array.clone_from_mapping(mapping))
+            }
+        }
+    }
+}
+
 #[cfg(all(feature = "memmap", target_endian = "little"))]
 impl From<MmapArray> for StorageViewWrap {
     fn from(s: MmapArray) -> Self {

src/chunks/vocab/subword.rs

@@ -357,7 +357,10 @@ where
     /// N-grams in the range `(self.min_n..self.max_n)` are extracted from the words in the
     /// vocabulary, each of these gets assigned an index from the `BucketIndexer` which is used to
     /// determine the index in the explicit subword vocab.
-    pub fn to_explicit(&self) -> ExplicitSubwordVocab {
+    ///
+    /// The second item in the returned tuple holds the `bucket -> explicit_index` mapping for
+    /// all buckets hit by `self`.
+    pub fn to_explicit(&self) -> (ExplicitSubwordVocab, HashMap<u64, usize>) {
         let mut ngram_index = HashMap::new();
         let SubwordVocab {
             words,
@@ -375,8 +378,11 @@ where
                 ngram_index.entry(ngram.into()).or_insert(idx);
             }
         }
-        let indexer = ExplicitIndexer::new_with_indices(ngram_index);
-        ExplicitSubwordVocab::new(words.to_owned(), *min_n, *max_n, indexer)
+        let (indexer, mapping) = ExplicitIndexer::new_with_indices(ngram_index);
+        (
+            ExplicitSubwordVocab::new(words.to_owned(), *min_n, *max_n, indexer),
+            mapping,
+        )
     }
 }
 
@@ -407,7 +413,7 @@ impl SubwordVocab<ExplicitIndexer> {
 
         let words = read_vocab_items(read, words_len as usize)?;
         let ngrams = read_ngrams_with_indices(read, ngrams_len as usize)?;
-        let indexer = ExplicitIndexer::new_with_indices(ngrams);
+        let (indexer, _) = ExplicitIndexer::new_with_indices(ngrams);
         Ok(SubwordVocab::new(words, min_n, max_n, indexer))
     }
 
@@ -560,15 +566,15 @@ mod tests {
             ("<t".to_owned(), 2),
         ];
 
-        ExplicitSubwordVocab::new(words, 2, 3, ExplicitIndexer::new_with_indices(ngrams))
+        ExplicitSubwordVocab::new(words, 2, 3, ExplicitIndexer::new_with_indices(ngrams).0)
     }
 
     #[test]
     fn test_conversion() {
         let words = vec!["groß".to_owned(), "allerdings".to_owned()];
         let indexer = FinalfusionHashIndexer::new(21);
         let bucket_vocab = SubwordVocab::new(words, 3, 6, indexer);
-        let explicit = bucket_vocab.to_explicit();
+        let (explicit, _) = bucket_vocab.to_explicit();
         let dings = StrWithCharLen::new("dings");
         let gro = StrWithCharLen::new("<gro");
         let dings_expl_idx = explicit.indexer().index_ngram(&dings);

src/compat/text.rs

@@ -265,10 +265,11 @@ where
 
     let matrix = Array2::from_shape_vec(shape, data).map_err(Error::Shape)?;
 
-    Ok(Embeddings::new_without_norms(
+    Ok(Embeddings::new_with_maybe_norms(
         None,
         SimpleVocab::new(words),
         NdArray::new(matrix),
+        None,
     ))
 }
 

src/compat/word2vec.rs

@@ -107,10 +107,11 @@ where
                 .map_err(|e| Error::io_error("Cannot read word embedding", e))?;
         }
 
-        Ok(Embeddings::new_without_norms(
+        Ok(Embeddings::new_with_maybe_norms(
             None,
             SimpleVocab::new(words),
             NdArray::new(matrix),
+            None,
         ))
     }
 }