kalosm-language-model 0.3.3

A common interface for language models/transformers
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

use futures_util::Future;

use crate::{Embedder, EmbedderExt, Embedding, VectorSpace};

/// Convert a type into an embedding with an embedding model.
pub trait IntoEmbedding<S: VectorSpace> {
    /// Convert the type into an embedding with the given embedding model.
    fn into_embedding<E: Embedder<VectorSpace = S>>(
        self,
        embedder: &E,
    ) -> impl Future<Output = anyhow::Result<Embedding<S>>>;

    /// Convert the type into a query embedding with the given embedding model.
    fn into_query_embedding<E: Embedder<VectorSpace = S>>(
        self,
        embedder: &E,
    ) -> impl Future<Output = anyhow::Result<Embedding<S>>>;
}

/// Convert any type that implements [`ToString`] into an embedding with an embedding model.
impl<S: ToString, V: VectorSpace> IntoEmbedding<V> for S {
    async fn into_embedding<E: Embedder<VectorSpace = V>>(
        self,
        embedder: &E,
    ) -> anyhow::Result<Embedding<V>> {
        embedder.embed(self).await
    }

    async fn into_query_embedding<E: Embedder<VectorSpace = V>>(
        self,
        embedder: &E,
    ) -> anyhow::Result<Embedding<V>> {
        embedder.embed_query(self).await
    }
}

/// Convert an embedding of the same vector space into an embedding with an embedding model.
impl<S: VectorSpace> IntoEmbedding<S> for Embedding<S> {
    async fn into_embedding<E: Embedder<VectorSpace = S>>(
        self,
        _: &E,
    ) -> anyhow::Result<Embedding<S>> {
        Ok(self)
    }

    async fn into_query_embedding<E: Embedder<VectorSpace = S>>(
        self,
        _: &E,
    ) -> anyhow::Result<Embedding<S>> {
        Ok(self)
    }
}