tf_idf_vectorizer/vectorizer/
serde.rs

1use std::sync::Arc;
2use std::hash::Hash;
3
4use ahash::RandomState;
5use num_traits::Num;
6use serde::{ser::SerializeStruct, Deserialize, Serialize};
7
8use crate::{Corpus, TFIDFVectorizer, utils::datastruct::{map::IndexMap, vector::ZeroSpVecTrait}, vectorizer::{IDFVector, KeyRc, TFVector, tfidf::{DefaultTFIDFEngine, TFIDFEngine}}};
9
10/// Data structure for deserializing TFIDFVectorizer.
11/// This struct does not contain references, so it can be serialized.
12/// Use the `into_tf_idf_vectorizer` method to convert to `TFIDFVectorizer`.
13#[derive(Debug, Deserialize, Serialize)]
14pub struct TFIDFData<N = f32, K = String, E = DefaultTFIDFEngine>
15where
16    N: Num + Copy,
17    E: TFIDFEngine<N, K>,
18    K: Clone + Eq + Hash,
19{
20    /// TF vectors for documents
21    pub documents: IndexMap<KeyRc<K>, TFVector<N>>,
22    /// Token dimension sample for TF vectors
23    pub token_dim_sample: Vec<Box<str>>,
24    /// IDF vector
25    #[serde(default, skip_serializing, skip_deserializing)]
26    pub idf: Option<IDFVector<N>>,
27    #[serde(default, skip_serializing, skip_deserializing)]
28    _marker: std::marker::PhantomData<E>,
29}
30
31impl<N, K, E> TFIDFData<N, K, E>
32where
33    N: Num + Copy + Into<f64> + Send + Sync,
34    E: TFIDFEngine<N, K>,
35    K: Clone + Send + Sync + Eq + Hash,
36{
37    /// Convert `TFIDFData` into `TFIDFVectorizer`.
38    /// `corpus_ref` is a reference to the corpus.
39    pub fn into_tf_idf_vectorizer(self, corpus_ref: Arc<Corpus>) -> TFIDFVectorizer<N, K, E>
40    {
41        let raw_iter = self.documents.iter();
42        let mut token_dim_rev_index: IndexMap<Box<str>, Vec<KeyRc<K>>, RandomState> =
43            IndexMap::with_capacity(self.token_dim_sample.len());
44        self.token_dim_sample.iter().for_each(|token| {
45            token_dim_rev_index.insert(&token.clone(), Vec::new());
46        });
47        for (key, doc) in raw_iter {
48            doc.tf_vec.raw_iter().for_each(|(idx, _)| {
49                let token = &self.token_dim_sample[idx];
50                token_dim_rev_index
51                    .get_mut(token).unwrap()
52                    .push(key.clone());
53            });
54        }
55
56        let mut instance = TFIDFVectorizer {
57            documents: self.documents,
58            token_dim_rev_index: token_dim_rev_index,
59            corpus_ref,
60            idf_cache: IDFVector::new(),
61            _marker: std::marker::PhantomData,
62        };
63        instance.update_idf();
64        instance
65    }
66}
67
68impl<N, K, E> Serialize for TFIDFVectorizer<N, K, E>
69where
70    N: Num + Copy + Serialize + Into<f64> + Send + Sync,
71    K: Serialize + Clone + Send + Sync + Eq + Hash,
72    E: TFIDFEngine<N, K>,
73{
74    /// Serialize TFIDFVectorizer.
75    /// This struct contains references, so they are excluded from serialization.
76    /// Use `TFIDFData` for deserialization.
77    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
78    where
79        S: serde::Serializer,
80    {
81        let mut state = serializer.serialize_struct("TFIDFVectorizer", 2)?;
82        state.serialize_field("documents", &self.documents)?;
83        state.serialize_field("token_dim_sample", &self.token_dim_rev_index.keys())?;
84        state.end()
85    }
86}