Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[ENH] Create
RoboflowEmbeddingFunction
(chroma-core#1434)
## Description of changes This PR adds a new `RoboflowEmbeddingFunction` with which a user can calculate CLIP text and image embeddings using [Roboflow Inference](https://inference.roboflow.com). ## Test plan You can test the embedding function using the following code: ```python import chromadb import os from chromadb.utils.embedding_functions import RoboflowEmbeddingFunction import uuid from PIL import Image client = chromadb.PersistentClient(path="database") collection = client.create_collection(name="images", metadata={"hnsw:space": "cosine"}) # collection = client.get_collection(name="images") IMAGE_DIR = "images/train/images/" SERVER_URL = "https://infer.roboflow.com" API_KEY = "" results = [] ef = RoboflowEmbeddingFunction(API_KEY) documents = [os.path.join(IMAGE_DIR, img) for img in os.listdir(IMAGE_DIR)] embeddings = ef(images = [img for img in documents]) ids = [str(uuid.uuid4()) for _ in range(len(documents))] print(len(embeddings)) collection.add( embeddings=embeddings, documents=documents, ids=ids, ) query = ef(prompt = "baseball") results = collection.query( query_embeddings=query, n_results=3 ) top_result = results["documents"] for i in top_result: print(i) ``` You will need a [Roboflow API key](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key) - [X] Tests pass locally with `pytest` for python, `yarn test` for js ## Documentation Changes I will file a PR to the `chroma-core/docs` repository with documentation. --------- Co-authored-by: Anton Troynikov <[email protected]>
- Loading branch information