abu-rag 0.2.0

Retrieval-Augmented Generation system
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use super::Vector;
use num_traits::Zero;
use num_traits::Float;

pub trait VectorMetric : Send + Sync {
    fn score<V: Vector>(a: V, b: V) -> Result<V::Scalar, VectorMetricError>;
}

#[derive(Debug, thiserror::Error)]
pub enum VectorMetricError {
    #[error("vector len no equal {0} != {1}")]
    VectorLen(usize, usize),
}

// ============================================================================ //
//                 cosine similarity
// ============================================================================ //

pub struct Cosine;

impl Cosine {
    pub fn cosine_similarity<V: Vector>(a: V, b: V) -> Result<V::Scalar, VectorMetricError> {
        if a.len() != b.len() {
            return Err(VectorMetricError::VectorLen(a.len(), b.len()))
        }
    
        let mut dot_product = V::Scalar::zero();
        let mut norm_a = V::Scalar::zero();
        let mut norm_b = V::Scalar::zero();
    
        for (x, y) in a.into_iter().zip(b.into_iter()) {
            dot_product += x * y;
            norm_a += x * x;
            norm_b += y * y;
        }
    
        let score = if norm_a == V::Scalar::zero() || norm_b == V::Scalar::zero() {
            V::Scalar::zero()
        } else {
            dot_product / (norm_a.sqrt() * norm_b.sqrt())
        };

        Ok(score)
    }
}

impl VectorMetric for Cosine {
    #[inline]
    fn score<V: Vector>(a: V, b: V) -> Result<V::Scalar, VectorMetricError> {
        Self::cosine_similarity(a, b)
    }
}

// ============================================================================ //
//                 cosine similarity
// ============================================================================ //

pub struct L2;

impl L2 {
    pub fn l2_distance<V: Vector>(a: V, b: V) -> Result<V::Scalar, VectorMetricError> {
        if a.len() != b.len() {
            return Err(VectorMetricError::VectorLen(a.len(), b.len()))
        }

        let score = a.into_iter().zip(b.into_iter())
            .map(|(a, b)| (a - b).powi(2))
            .sum::<V::Scalar>()
            .sqrt();

        Ok(-score)
    }
}

impl VectorMetric for L2 {
    #[inline]
    fn score<V: Vector>(a: V, b: V) -> Result<V::Scalar, VectorMetricError> {
        Self::l2_distance(a, b)
    }
}

// ============================================================================ //
//                 cosine similarity
// ============================================================================ //

pub struct Dot;

impl Dot {
    pub fn dot<V: Vector>(a: V, b: V) -> Result<V::Scalar, VectorMetricError> {
        if a.len() != b.len() {
            return Err(VectorMetricError::VectorLen(a.len(), b.len()))
        }

        let score = a.into_iter().zip(b.into_iter())
            .map(|(x, y)| x * y)
            .sum();

        Ok(score)
    }
}

impl VectorMetric for Dot {
    #[inline]
    fn score<V: Vector>(a: V, b: V) -> Result<V::Scalar, VectorMetricError> {
        Self::dot(a, b)
    }
}