word_vec_rs/
vector.rs

1use crate::as_vector::AsVectorRef;
2use borrowme::borrowme;
3use nalgebra::DVectorView;
4use std::ops::Add;
5
6/// A single WordVector
7#[borrowme]
8#[borrowed_attr(derive(Copy))]
9#[derive(Debug, Clone, PartialEq)]
10#[cfg_attr(feature = "serde", derive(serde::Serialize))]
11#[owned_attr(cfg_attr(feature = "serde", derive(serde::Deserialize)))]
12pub struct Vector<'v, 't> {
13    #[borrowme(borrow_with = Vec::as_slice)]
14    data: &'v [f32],
15    term: &'t str,
16}
17
18impl<'v, 't> Vector<'v, 't> {
19    #[inline]
20    pub fn new(data: &'v [f32], term: &'t str) -> Self {
21        Self { data, term }
22    }
23
24    #[inline]
25    pub fn data(&self) -> &[f32] {
26        self.data
27    }
28
29    #[inline]
30    pub fn term(&self) -> &str {
31        self.term
32    }
33
34    #[inline]
35    pub fn dim(&self) -> usize {
36        self.data.len()
37    }
38
39    /// Calculates the cosine similarity between two words.
40    pub fn cosine<'v2, 't2, R>(&self, other: &R) -> f32
41    where
42        R: AsVectorRef<'v2, 't2>,
43    {
44        let other = other.as_vec_ref();
45
46        let dot = self.dot(&other);
47        if dot == 0.0 {
48            return 0.0;
49        }
50
51        let div = self.length() * other.length();
52        if div == 0.0 {
53            return 0.0;
54        }
55
56        dot / div
57    }
58
59    /// Calculates the dot product of two vectors
60    pub fn dot<'v2, 't2, R>(&self, other: &R) -> f32
61    where
62        R: AsVectorRef<'v2, 't2>,
63    {
64        // self.vec().dot(&other.as_vec_ref().vec())
65        let other = other.as_vec_ref();
66        self.data
67            .iter()
68            .zip(other.data.iter())
69            .map(|(a, b)| a * b)
70            .sum()
71    }
72
73    #[inline]
74    pub fn vec(&self) -> DVectorView<'_, f32> {
75        DVectorView::from_slice(self.data, 1)
76    }
77
78    /// Calculates the 2-norm
79    #[inline]
80    pub fn length(&self) -> f32 {
81        // self.vec().norm()
82        self.data.iter().map(|i| i.powi(2)).sum::<f32>().sqrt()
83    }
84}
85
86impl OwnedVector {
87    #[inline]
88    pub fn new_raw(data: Vec<f32>, term: String) -> Self {
89        Self { data, term }
90    }
91
92    #[inline]
93    pub fn new(data: &[f32], term: &str) -> Self {
94        borrowme::ToOwned::to_owned(&Vector::new(data, term))
95    }
96
97    /// Returns a reference to the data of the owned vector.
98    #[inline]
99    pub fn as_ref(&self) -> Vector {
100        Vector::new(&self.data, &self.term)
101    }
102
103    #[inline]
104    pub fn data(&self) -> &[f32] {
105        &self.data
106    }
107
108    #[inline]
109    pub fn term(&self) -> &str {
110        &self.term
111    }
112
113    #[inline]
114    pub fn dim(&self) -> usize {
115        self.data.len()
116    }
117
118    /// Calculates the 2-norm
119    #[inline]
120    pub fn length(&self) -> f32 {
121        self.as_ref().length()
122    }
123
124    /// Calculates the 2-norm
125    #[inline]
126    pub fn dot<'v, 't, R: AsVectorRef<'v, 't>>(&self, other: &R) -> f32 {
127        self.as_ref().dot(other)
128    }
129
130    /// Calculates the cosine similarity between two vectors.
131    #[inline]
132    pub fn cosine<'v2, 't2, R>(&self, other: &R) -> f32
133    where
134        R: AsVectorRef<'v2, 't2>,
135    {
136        self.as_ref().cosine(other)
137    }
138}
139
140impl<'v, 't, 'v2, 't2, T> Add<T> for Vector<'v, 't>
141where
142    T: AsVectorRef<'v2, 't2>,
143{
144    type Output = OwnedVector;
145
146    fn add(self, rhs: T) -> Self::Output {
147        let rhs = rhs.as_vec_ref();
148        assert_eq!(self.dim(), rhs.dim());
149
150        let data: Vec<_> = self
151            .data
152            .iter()
153            .zip(rhs.data.iter())
154            .map(|i| i.0 + i.1)
155            .collect();
156
157        OwnedVector::new_raw(data, format!("{} {}", self.term, rhs.term()))
158    }
159}
160
161impl<'v, 't, T> Add<T> for OwnedVector
162where
163    T: AsVectorRef<'v, 't>,
164{
165    type Output = OwnedVector;
166
167    fn add(self, rhs: T) -> Self::Output {
168        self.as_ref() + rhs
169    }
170}
171
172impl<'v, 't> AsVectorRef<'v, 't> for &Vector<'v, 't> {
173    #[inline]
174    fn as_vec_ref(&self) -> Vector<'v, 't> {
175        **self
176    }
177}
178
179impl<'v, 't> AsVectorRef<'v, 't> for Vector<'v, 't> {
180    #[inline]
181    fn as_vec_ref(&self) -> Vector<'v, 't> {
182        *self
183    }
184}
185
186impl<'a> AsVectorRef<'a, 'a> for &'a OwnedVector {
187    #[inline]
188    fn as_vec_ref(&self) -> Vector<'a, 'a> {
189        self.as_ref()
190    }
191}