OpenAI and Huggingface Providers (#16)

Providers for creating embeddings with OpenAI and Hugginface.

Includes
- base class
- simple integration testing
- async methods for clients that support it.

Author: Sam Partee

.gitignore

@@ -1,4 +1,4 @@
 __pycache__/
 redisvl.egg-info/
-.coverage
-scratch
+.coverage*
+scratch

conftest.py

@@ -37,3 +37,8 @@ def df():
         }
     )
     return data
+
+
+@pytest.fixture
+def openai_key():
+    return os.getenv("OPENAI_KEY")

redisvl/providers/__init__.py

@@ -0,0 +1,9 @@
+
+from redisvl.providers.openai import OpenAIProvider
+from redisvl.providers.huggingface import HuggingfaceProvider
+
+
+__all__ = [
+    "OpenAIProvider",
+    "HuggingfaceProvider"
+]

redisvl/providers/base.py

@@ -0,0 +1,43 @@
+from typing import Callable, Dict, List, Optional
+
+
+class BaseProvider:
+    def __init__(self, model: str, dims: int, api_config: Optional[Dict] = None):
+        self._dims = dims
+        self._model = model
+
+    @property
+    def model(self) -> str:
+        return self._model
+
+    @property
+    def dims(self) -> int:
+        return self._dims
+
+    def set_model(self, model, dims: int = None):
+        self._model = model
+        if dims:
+            self._dims = dims
+
+    def embed_many(
+        self, inputs: List[str], preprocess: callable = None, chunk_size: int = 1000
+    ) -> List[float]:
+        raise NotImplementedError
+
+    def embed(self, emb_input: str, preprocess: callable = None) -> List[float]:
+        raise NotImplementedError
+
+    async def aembed_many(
+        self, inputs: List[str], preprocess: callable = None, chunk_size: int = 1000
+    ) -> List[float]:
+        raise NotImplementedError
+
+    async def aembed(self, emb_input: str, preprocess: callable = None) -> List[float]:
+        raise NotImplementedError
+
+    def batchify(self, seq: list, size: int, preprocess: Callable = None):
+        for pos in range(0, len(seq), size):
+            if preprocess:
+                yield [preprocess(chunk) for chunk in seq[pos : pos + size]]
+            else:
+                yield [chunk for chunk in seq[pos : pos + size]]

redisvl/providers/huggingface.py

@@ -0,0 +1,29 @@
+from typing import Dict, List, Optional
+
+from sentence_transformers import SentenceTransformer
+
+from redisvl.providers.base import BaseProvider
+
+
+class HuggingfaceProvider(BaseProvider):
+    def __init__(self, model: str, api_config: Optional[Dict] = None):
+        # TODO set dims based on model
+        dims = 768
+        super().__init__(model, dims, api_config)
+        self._model_client = SentenceTransformer(model)
+
+    def embed(self, emb_input: str, preprocess: callable = None) -> List[float]:
+        if preprocess:
+            emb_input = preprocess(emb_input)
+        embedding = self._model_client.encode([emb_input])[0]
+        return embedding.tolist()
+
+    def embed_many(
+        self, inputs: List[str], preprocess: callable = None, chunk_size: int = 1000
+    ) -> List[List[float]]:
+
+        embeddings = []
+        for batch in self.batchify(inputs, chunk_size, preprocess):
+            batch_embeddings = self._model_client.encode(batch)
+            embeddings.extend([embedding.tolist() for embedding in batch_embeddings])
+        return embeddings

redisvl/providers/openai.py

@@ -0,0 +1,48 @@
+from typing import Dict, List, Optional
+
+import openai
+
+from redisvl.providers.base import BaseProvider
+
+
+class OpenAIProvider(BaseProvider):
+    def __init__(self, model: str, api_config: Optional[Dict] = None):
+        dims = 1536
+        super().__init__(model, dims, api_config)
+        if not api_config:
+            raise ValueError("OpenAI API key is required in api_config")
+
+        openai.api_key = api_config.get("api_key", None)
+        self._model_client = openai.Embedding
+
+    def embed_many(
+        self, inputs: List[str], preprocess: callable = None, chunk_size: int = 1000
+    ) -> List[List[float]]:
+
+        results = []
+        for batch in self.batchify(inputs, chunk_size, preprocess):
+            response = self._model_client.create(input=batch, engine=self._model)
+            results += [r["embedding"] for r in response["data"]]
+        return results
+
+    def embed(self, emb_input: str, preprocess: callable = None) -> List[float]:
+        if preprocess:
+            emb_input = preprocess(emb_input)
+        result = self._model_client.create(input=[emb_input], engine=self._model)
+        return result["data"][0]["embedding"]
+
+    async def aembed_many(
+        self, inputs: List[str], preprocess: callable = None, chunk_size: int = 1000
+    ) -> List[List[float]]:
+
+        results = []
+        for batch in self.batchify(inputs, chunk_size, preprocess):
+            response = await self._model_client.acreate(input=batch, engine=self._model)
+            results += [r["embedding"] for r in response["data"]]
+        return results
+
+    async def aembed(self, emb_input: str, preprocess: callable = None) -> List[float]:
+        if preprocess:
+            emb_input = preprocess(emb_input)
+        result = await self._model_client.acreate(input=[emb_input], engine=self._model)
+        return result["data"][0]["embedding"]

tests/integration/test_providers.py

@@ -0,0 +1,57 @@
+import pytest
+from redisvl.providers import (
+    HuggingfaceProvider,
+    OpenAIProvider
+)
+
+
+@pytest.fixture(params=[HuggingfaceProvider, OpenAIProvider])
+def provider(request, openai_key):
+    # Here we use actual models for integration test
+    if request.param == HuggingfaceProvider:
+        return request.param(model="sentence-transformers/all-mpnet-base-v2")
+    elif request.param == OpenAIProvider:
+        return request.param(model="text-embedding-ada-002", api_config={
+            "api_key": openai_key
+            })
+
+def test_provider_embed(provider):
+    text = 'This is a test sentence.'
+    embedding = provider.embed(text)
+
+    assert isinstance(embedding, list)
+    assert len(embedding) == provider.dims
+
+def test_provider_embed_many(provider):
+    texts = ['This is the first test sentence.', 'This is the second test sentence.']
+    embeddings = provider.embed_many(texts)
+
+    assert isinstance(embeddings, list)
+    assert len(embeddings) == len(texts)
+    assert all(isinstance(emb, list) and len(emb) == provider.dims for emb in embeddings)
+
+
+@pytest.fixture(params=[OpenAIProvider])
+def aprovider(request, openai_key):
+    # Here we use actual models for integration test
+    if request.param == OpenAIProvider:
+        return request.param(model="text-embedding-ada-002", api_config={
+            "api_key": openai_key
+            })
+
+@pytest.mark.asyncio
+async def test_provider_aembed(aprovider):
+    text = 'This is a test sentence.'
+    embedding = await aprovider.aembed(text)
+
+    assert isinstance(embedding, list)
+    assert len(embedding) == aprovider.dims
+
+@pytest.mark.asyncio
+async def test_provider_aembed_many(aprovider):
+    texts = ['This is the first test sentence.', 'This is the second test sentence.']
+    embeddings = await aprovider.aembed_many(texts)
+
+    assert isinstance(embeddings, list)
+    assert len(embeddings) == len(texts)
+    assert all(isinstance(emb, list) and len(emb) == aprovider.dims for emb in embeddings)