Skip to content

dspy.KNN

dspy.KNN(k: int, trainset: list[Example], vectorizer: Embedder)

一个k近邻检索器，用于从训练集中查找相似的示例。

参数:

名称 类型 描述 默认值
k int

要检索的最近邻数量

 必填
trainset list[Example]

用于搜索的训练示例列表

 必填
vectorizer Embedder

用于向量化的 Embedder

 必填
示例 
import dspy
from sentence_transformers import SentenceTransformer

# Create a training dataset with examples
trainset = [
    dspy.Example(input="hello", output="world"),
    # ... more examples ...
]

# Initialize KNN with a sentence transformer model
knn = KNN(
    k=3,
    trainset=trainset,
    vectorizer=dspy.Embedder(SentenceTransformer("all-MiniLM-L6-v2").encode)
)

# Find similar examples
similar_examples = knn(input="hello")
Source code in dspy/predict/knn.py 
 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
def __init__(self, k: int, trainset: list[Example], vectorizer: Embedder):
    """
    A k-nearest neighbors retriever that finds similar examples from a training set.

    Args:
        k: Number of nearest neighbors to retrieve
        trainset: List of training examples to search through
        vectorizer: The `Embedder` to use for vectorization

    Example:
        ```python
        import dspy
        from sentence_transformers import SentenceTransformer

        # Create a training dataset with examples
        trainset = [
            dspy.Example(input="hello", output="world"),
            # ... more examples ...
        ]

        # Initialize KNN with a sentence transformer model
        knn = KNN(
            k=3,
            trainset=trainset,
            vectorizer=dspy.Embedder(SentenceTransformer("all-MiniLM-L6-v2").encode)
        )

        # Find similar examples
        similar_examples = knn(input="hello")
        ```
    """
    self.k = k
    self.trainset = trainset
    self.embedding = vectorizer
    trainset_casted_to_vectorize = [
        " | ".join([f"{key}: {value}" for key, value in example.items() if key in example._input_keys])
        for example in self.trainset
    ]
    self.trainset_vectors = self.embedding(trainset_casted_to_vectorize).astype(np.float32)

函数

__call__(**kwargs) -> list

Source code in dspy/predict/knn.py 
48
49
50
51
52
def __call__(self, **kwargs) -> list:
    input_example_vector = self.embedding([" | ".join([f"{key}: {val}" for key, val in kwargs.items()])])
    scores = np.dot(self.trainset_vectors, input_example_vector.T).squeeze()
    nearest_samples_idxs = scores.argsort()[-self.k :][::-1]
    return [self.trainset[cur_idx] for cur_idx in nearest_samples_idxs]

:::