Repository: stanfordnlp/dspy
Pull request #930 (2024-04-29 18:35:51)

Clean up knn related functionalities a bit
Clean up `predict/knn` and `teleprompt/knn_fewshot`:
* Added a optional `vectorizer` instead of using hardcoded `SentenceTransformer` one
* Clean up constructor a bit


Patch:
diff --git a/docs/docs/tutorials/other_tutorial.md b/docs/docs/tutorials/other_tutorial.md
index c37caed1..d276a5db 100644
--- a/docs/docs/tutorials/other_tutorial.md
+++ b/docs/docs/tutorials/other_tutorial.md
@@ -24,5 +24,4 @@ sidebar_position: 99999
 - [DSPy webinar with MLOps Learners](https://www.youtube.com/watch?v=im7bCLW2aM4), a bit longer with Q&A.
 - Hands-on Overviews of DSPy by the community: [DSPy Explained! by Connor Shorten](https://www.youtube.com/watch?v=41EfOY0Ldkc), [DSPy explained by code_your_own_ai](https://www.youtube.com/watch?v=ycfnKPxBMck), [DSPy Crash Course by AI Bites](https://youtu.be/5-zgASQKkKQ?si=3gnmVouT5_rpk_nu)
 - Interviews: [Weaviate Podcast in-person](https://www.youtube.com/watch?v=CDung1LnLbY), and you can find 6-7 other remote podcasts on YouTube from a few different perspectives/audiences.
-- **Tracing in DSPy** with Arize Phoenix: [Tutorial for tracing your prompts and the steps of your DSPy programs](https://colab.research.google.com/github/Arize-ai/phoenix/blob/main/tutorials/tracing/dspy_tracing_tutorial.ipynb)
-- **Tracing & Optimization Tracking in DSPy** with Parea AI: [Tutorial on tracing & evaluating a DSPy RAG program](https://docs.parea.ai/tutorials/dspy-rag-trace-evaluate/tutorial)
\ No newline at end of file
+- **Tracing in DSPy** with Arize Phoenix: [Tutorial for tracing your prompts and the steps of your DSPy programs](https://colab.research.google.com/github/Arize-ai/phoenix/blob/main/tutorials/tracing/dspy_tracing_tutorial.ipynb)
\ No newline at end of file
diff --git a/dspy/predict/knn.py b/dspy/predict/knn.py
index e275ccb6..1ce85fdb 100644
--- a/dspy/predict/knn.py
+++ b/dspy/predict/knn.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 
@@ -6,17 +6,20 @@ import dsp
 
 
 class KNN:
-    def __init__(self, k: int, trainset: List[dsp.Example]):
+    def __init__(self, k: int, trainset: List[dsp.Example], vectorizer: Optional[dsp.BaseSentenceVectorizer] = None):
         self.k = k
         self.trainset = trainset
-        self.vectorizer = dsp.SentenceTransformersVectorizer()
-        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.vectorizer = vectorizer or dsp.SentenceTransformersVectorizer()
+        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.vectorizer(trainset_casted_to_vectorize).astype(np.float32)
 
     def __call__(self, **kwargs) -> List[dsp.Example]:
         with dsp.settings.context(vectorizer=self.vectorizer):
             input_example_vector = self.vectorizer([" | ".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]
+            nearest_samples_idxs = scores.argsort()[-self.k :][::-1]
             train_sampled = [self.trainset[cur_idx] for cur_idx in nearest_samples_idxs]
             return train_sampled
diff --git a/dspy/teleprompt/knn_fewshot.py b/dspy/teleprompt/knn_fewshot.py
index ca940526..7211d913 100644
--- a/dspy/teleprompt/knn_fewshot.py
+++ b/dspy/teleprompt/knn_fewshot.py
@@ -1,17 +1,18 @@
 import types
-from typing import List
+from typing import List, Optional
 
 import dsp
+from dspy.predict.knn import KNN
 from dspy.teleprompt import BootstrapFewShot
 
 from .teleprompt import Teleprompter
 
 
 class KNNFewShot(Teleprompter):
-    def __init__(self, KNN, k: int, trainset: List[dsp.Example]):
-        self.KNN = KNN(k, trainset)
+    def __init__(self, k: int, trainset: List[dsp.Example], vectorizer: Optional[dsp.BaseSentenceVectorizer] = None):
+        self.KNN = KNN(k, trainset, vectorizer=vectorizer)
 
-    def compile(self, student, *, teacher=None, trainset, valset=None):
+    def compile(self, student, *, teacher=None, trainset=None, valset=None):
         student_copy = student.reset_copy()
 
         def forward_pass(*args, **kwargs):
diff --git a/examples/knn.ipynb b/examples/knn.ipynb
index 51a50654..77052db6 100644
--- a/examples/knn.ipynb
+++ b/examples/knn.ipynb
@@ -135,9 +135,8 @@
    "outputs": [],
    "source": [
     "from dspy.teleprompt import KNNFewShot\n",
-    "from dspy.predict.knn import KNN\n",
     "\n",
-    "knn_teleprompter = KNNFewShot(KNN, 7, trainset)\n",
+    "knn_teleprompter = KNNFewShot(7, trainset)\n",
     "compiled_knn = knn_teleprompter.compile(BasicQABot(), trainset=trainset)"
    ]
   },


Test patch:
diff --git a/tests/teleprompt/test_knn_fewshot.py b/tests/teleprompt/test_knn_fewshot.py
index b267d3dc..b0a9bf2e 100644
--- a/tests/teleprompt/test_knn_fewshot.py
+++ b/tests/teleprompt/test_knn_fewshot.py
@@ -1,6 +1,5 @@
 import pytest
 import dsp, dspy
-from dspy.predict.knn import KNN
 from dspy.teleprompt.knn_fewshot import KNNFewShot
 from dspy.utils.dummies import DummyLM, DummyVectorizer
 
@@ -19,7 +18,7 @@ def setup_knn_few_shot():
         mock_example("What is 2+2?", "4"),
     ]
     dsp.SentenceTransformersVectorizer = DummyVectorizer
-    knn_few_shot = KNNFewShot(KNN, k=2, trainset=trainset)
+    knn_few_shot = KNNFewShot(k=2, trainset=trainset)
     return knn_few_shot
 
 
@@ -60,9 +59,8 @@ def _test_knn_few_shot_compile(setup_knn_few_shot):
     assert len(compiled_student.predictor.demos) == 1
     assert compiled_student.predictor.demos[0].input == trainset[0].input
     assert compiled_student.predictor.demos[0].output == trainset[0].output
-    
     # Simulate a query that is similar to one of the training examples
-    output = compiled_student.forward(input = "What is the capital of Spain?").output
+    output = compiled_student.forward(input="What is the capital of Spain?").output
 
     print("CONVO")
     print(lm.get_convo(-1))