Add dspy.Embedding

[chenmoneygithub]

Very simple dspy.Embedding supports:

• Hosted embedding model, we use litellm to send the request, and only take the embedding fields from the response.
• Custom embedding model, users can pass a custom callable to dspy.Embedding, and the output is just the custom callable's output.

Added unit test for both scenarios.

Confirmed that the cache works:

(dspy) *[dspy-embedding][~/Documents/mlflow_team/dspy]$ python3
Python 3.12.4 | packaged by Anaconda, Inc. | (main, Jun 18 2024, 10:07:17) [Clang 14.0.6 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import dspy;import time;embedder = dspy.Embedding("openai/text-embedding-ada-002")
>>> start=time.time();geez = embedder(["oh my god", "holy geese and elephant", "ohhhh hoover dam"]);print("TOTAL TIME: ", time.time()-start)
TOTAL TIME:  0.963151216506958
>>> start=time.time();geez = embedder(["oh my god", "holy geese and elephant", "ohhhh hoover dam"]);print("TOTAL TIME: ", time.time()-start)
TOTAL TIME:  0.009987115859985352
>>> start=time.time();geez = embedder(["oh my god", "holy geese and elephant", "ohhhh hoover dam"]);print("TOTAL TIME: ", time.time()-start)
TOTAL TIME:  0.010275125503540039

[okhat]

I think this needs to be done before LiteLLM is imported anywhere in DSPy, for it to have an effect?

[chenmoneygithub]

I searched their code, and this env var is read at runtime: https://github.com/BerriAI/litellm/blob/5652c375b3e22bab6704e93058c868620c72d6ee/litellm/__init__.py#L309, so our current order should be okay.

[okhat]

Can we ensure the output of this is a numpy tensor or something? Both for litellm and for callables.

[chenmoneygithub]

sounds good!

[okhat]

An ideal version of this PR would involve improving the docs at this page:

https://dspy-docs.vercel.app/quick-start/getting-started-02/ (see the second cell, or see below)

import torch
import functools
from litellm import embedding as Embed

with open("test_collection.jsonl") as f:
    corpus = [ujson.loads(line) for line in f]

index = torch.load('index.pt', weights_only=True)
max_characters = 4000 # >98th percentile of document lengths

@functools.lru_cache(maxsize=None)
def search(query, k=5):
    query_embedding = torch.tensor(Embed(input=query, model="text-embedding-3-small").data[0]['embedding'])
    topk_scores, topk_indices = torch.matmul(index, query_embedding).topk(k)
    topK = [dict(score=score.item(), **corpus[idx]) for idx, score in zip(topk_indices, topk_scores)]
    return [doc['text'][:max_characters] for doc in topK]

I'd love to get the same functionality but without that complexity...