From 0bce1a9d7f06dc7d49d08aaa259a69f40736f72d Mon Sep 17 00:00:00 2001 From: chenmoneygithub Date: Thu, 31 Oct 2024 21:02:47 -0700 Subject: [PATCH 1/2] Add embedding model --- dspy/clients/__init__.py | 13 ++++++- dspy/clients/embedding.py | 41 +++++++++++++++++++++ dspy/clients/lm.py | 9 ----- tests/clients/test_embedding.py | 64 +++++++++++++++++++++++++++++++++ 4 files changed, 117 insertions(+), 10 deletions(-) create mode 100644 dspy/clients/embedding.py create mode 100644 tests/clients/test_embedding.py diff --git a/dspy/clients/__init__.py b/dspy/clients/__init__.py index 6a63509f5b..d1090b5285 100644 --- a/dspy/clients/__init__.py +++ b/dspy/clients/__init__.py @@ -1,2 +1,13 @@ from .lm import LM -from .base_lm import BaseLM \ No newline at end of file +from .base_lm import BaseLM +import litellm +import os +from pathlib import Path +from litellm.caching import Cache + +DISK_CACHE_DIR = os.environ.get("DSPY_CACHEDIR") or os.path.join(Path.home(), ".dspy_cache") +litellm.cache = Cache(disk_cache_dir=DISK_CACHE_DIR, type="disk") +litellm.telemetry = False + +if "LITELLM_LOCAL_MODEL_COST_MAP" not in os.environ: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" diff --git a/dspy/clients/embedding.py b/dspy/clients/embedding.py new file mode 100644 index 0000000000..3e2a881390 --- /dev/null +++ b/dspy/clients/embedding.py @@ -0,0 +1,41 @@ +import litellm + + +class Embedding: + """DSPy embedding class. + + The class for computing embeddings for text inputs. This class supports both hosted embedding + models like OpenAI's text-embedding-3-small as well as custom embedding models that are provided as a function. + When hosted embedding models are used, this class relies on litellm to call the embedding model. When a custom + embedding model is used, this class directly passes the inputs to the model function. + + Args: + model: The embedding model to use. This can be either a string (representing the name of the hosted embedding + model, must be an embedding model supported by litellm) or a callable that represents a custom embedding + model. + """ + + def __init__(self, model): + self.model = model + + def __call__(self, inputs, caching=True, **kwargs): + """Compute embeddings for the given inputs. + + Args: + inputs: The inputs to compute embeddings for, can be a single string or a list of strings. + caching: Whether to cache the embedding response, only valid when using a hosted embedding model. + kwargs: Additional keyword arguments to pass to the embedding model. + + Returns: + A list of embeddings, one for each input, in the same order as the inputs. Or the output of the custom + embedding model. + """ + if isinstance(inputs, str): + inputs = [inputs] + if isinstance(self.model, str): + embedding_response = litellm.embedding(model=self.model, input=inputs, caching=caching, **kwargs) + return [data["embedding"] for data in embedding_response.data] + elif callable(self.model): + return self.model(inputs, **kwargs) + else: + raise ValueError(f"`model` in `dspy.Embedding` must be a string or a callable, but got {type(self.model)}.") diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py index 22e37019f4..ce11f460b7 100644 --- a/dspy/clients/lm.py +++ b/dspy/clients/lm.py @@ -5,24 +5,15 @@ import uuid from concurrent.futures import ThreadPoolExecutor from datetime import datetime -from pathlib import Path from typing import Any, Dict, List, Literal, Optional import litellm import ujson -from litellm.caching import Cache from dspy.clients.finetune import FinetuneJob, TrainingMethod from dspy.clients.lm_finetune_utils import execute_finetune_job, get_provider_finetune_job_class from dspy.utils.callback import BaseCallback, with_callbacks -DISK_CACHE_DIR = os.environ.get("DSPY_CACHEDIR") or os.path.join(Path.home(), ".dspy_cache") -litellm.cache = Cache(disk_cache_dir=DISK_CACHE_DIR, type="disk") -litellm.telemetry = False - -if "LITELLM_LOCAL_MODEL_COST_MAP" not in os.environ: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - GLOBAL_HISTORY = [] logger = logging.getLogger(__name__) diff --git a/tests/clients/test_embedding.py b/tests/clients/test_embedding.py new file mode 100644 index 0000000000..f117619542 --- /dev/null +++ b/tests/clients/test_embedding.py @@ -0,0 +1,64 @@ +import pytest +from unittest.mock import Mock, patch +import numpy as np + +from dspy.clients.embedding import Embedding + + +# Mock response format similar to litellm's embedding response. +class MockEmbeddingResponse: + def __init__(self, embeddings): + self.data = [{"embedding": emb} for emb in embeddings] + self.usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + self.model = "mock_model" + self.object = "list" + + +def test_litellm_embedding(): + model = "text-embedding-ada-002" + inputs = ["hello", "world"] + mock_embeddings = [ + [0.1, 0.2, 0.3], # embedding for "hello" + [0.4, 0.5, 0.6], # embedding for "world" + ] + + with patch("litellm.embedding") as mock_litellm: + # Configure mock to return proper response format. + mock_litellm.return_value = MockEmbeddingResponse(mock_embeddings) + + # Create embedding instance and call it. + embedding = Embedding(model) + result = embedding(inputs) + + # Verify litellm was called with correct parameters. + mock_litellm.assert_called_once_with(model=model, input=inputs, caching=True) + + assert len(result) == len(inputs) + assert result == mock_embeddings + + +def test_callable_embedding(): + inputs = ["hello", "world", "test"] + + expected_embeddings = [ + [0.1, 0.2, 0.3], # embedding for "hello" + [0.4, 0.5, 0.6], # embedding for "world" + [0.7, 0.8, 0.9], # embedding for "test" + ] + + def mock_embedding_fn(texts): + # Simple callable that returns random embeddings. + return expected_embeddings + + # Create embedding instance with callable + embedding = Embedding(mock_embedding_fn) + result = embedding(inputs) + + assert result == expected_embeddings + + +def test_invalid_model_type(): + # Test that invalid model type raises ValueError + with pytest.raises(ValueError): + embedding = Embedding(123) # Invalid model type + embedding(["test"]) From 40bfd8ba5185aa423055ef9d398a9d78b9ce9366 Mon Sep 17 00:00:00 2001 From: chenmoneygithub Date: Fri, 1 Nov 2024 11:42:09 -0700 Subject: [PATCH 2/2] force return type to be numpy array --- dspy/clients/__init__.py | 1 + dspy/clients/embedding.py | 52 ++++++++++++++++++++++++++++----- tests/clients/test_embedding.py | 4 +-- 3 files changed, 47 insertions(+), 10 deletions(-) diff --git a/dspy/clients/__init__.py b/dspy/clients/__init__.py index d1090b5285..7f343152d0 100644 --- a/dspy/clients/__init__.py +++ b/dspy/clients/__init__.py @@ -1,5 +1,6 @@ from .lm import LM from .base_lm import BaseLM +from .embedding import Embedding import litellm import os from pathlib import Path diff --git a/dspy/clients/embedding.py b/dspy/clients/embedding.py index 3e2a881390..eec41c32b0 100644 --- a/dspy/clients/embedding.py +++ b/dspy/clients/embedding.py @@ -1,18 +1,55 @@ import litellm +import numpy as np class Embedding: """DSPy embedding class. - The class for computing embeddings for text inputs. This class supports both hosted embedding - models like OpenAI's text-embedding-3-small as well as custom embedding models that are provided as a function. - When hosted embedding models are used, this class relies on litellm to call the embedding model. When a custom - embedding model is used, this class directly passes the inputs to the model function. + The class for computing embeddings for text inputs. This class provides a unified interface for both: + + 1. Hosted embedding models (e.g. OpenAI's text-embedding-3-small) via litellm integration + 2. Custom embedding functions that you provide + + For hosted models, simply pass the model name as a string (e.g. "openai/text-embedding-3-small"). The class will use + litellm to handle the API calls and caching. + + For custom embedding models, pass a callable function that: + - Takes a list of strings as input. + - Returns embeddings as either: + - A 2D numpy array of float32 values + - A 2D list of float32 values + - Each row should represent one embedding vector Args: model: The embedding model to use. This can be either a string (representing the name of the hosted embedding model, must be an embedding model supported by litellm) or a callable that represents a custom embedding model. + + Examples: + Example 1: Using a hosted model. + + ```python + import dspy + + embedder = dspy.Embedding("openai/text-embedding-3-small") + embeddings = embedder(["hello", "world"]) + + assert embeddings.shape == (2, 1536) + ``` + + Example 2: Using a custom function. + + ```python + import dspy + + def my_embedder(texts): + return np.random.rand(len(texts), 10) + + embedder = dspy.Embedding(my_embedder) + embeddings = embedder(["hello", "world"]) + + assert embeddings.shape == (2, 10) + ``` """ def __init__(self, model): @@ -27,15 +64,14 @@ def __call__(self, inputs, caching=True, **kwargs): kwargs: Additional keyword arguments to pass to the embedding model. Returns: - A list of embeddings, one for each input, in the same order as the inputs. Or the output of the custom - embedding model. + A 2-D numpy array of embeddings, one embedding per row. """ if isinstance(inputs, str): inputs = [inputs] if isinstance(self.model, str): embedding_response = litellm.embedding(model=self.model, input=inputs, caching=caching, **kwargs) - return [data["embedding"] for data in embedding_response.data] + return np.array([data["embedding"] for data in embedding_response.data], dtype=np.float32) elif callable(self.model): - return self.model(inputs, **kwargs) + return np.array(self.model(inputs, **kwargs), dtype=np.float32) else: raise ValueError(f"`model` in `dspy.Embedding` must be a string or a callable, but got {type(self.model)}.") diff --git a/tests/clients/test_embedding.py b/tests/clients/test_embedding.py index f117619542..d12850e52c 100644 --- a/tests/clients/test_embedding.py +++ b/tests/clients/test_embedding.py @@ -34,7 +34,7 @@ def test_litellm_embedding(): mock_litellm.assert_called_once_with(model=model, input=inputs, caching=True) assert len(result) == len(inputs) - assert result == mock_embeddings + np.testing.assert_allclose(result, mock_embeddings) def test_callable_embedding(): @@ -54,7 +54,7 @@ def mock_embedding_fn(texts): embedding = Embedding(mock_embedding_fn) result = embedding(inputs) - assert result == expected_embeddings + np.testing.assert_allclose(result, expected_embeddings) def test_invalid_model_type():