openaivec 0.2.1__tar.gz

openaivec-0.2.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 anaregdesign
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

openaivec-0.2.1/PKG-INFO

@@ -0,0 +1,153 @@
+Metadata-Version: 2.1
+Name: openaivec
+Version: 0.2.1
+Summary: 
+Home-page: https://github.com/anaregdesign/vectorize-openai
+License: MIT
+Author: Hiroki Mizukami
+Author-email: hmizukami@microsoft.com
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: httpx[http2] (>=0.28.1,<0.29.0)
+Requires-Dist: openai (>=1.57.2,<2.0.0)
+Requires-Dist: pandas (>=2.2.3,<3.0.0)
+Requires-Dist: pyspark (>=3.5.1,<4.0.0)
+Project-URL: Repository, https://github.com/anaregdesign/vectorize-openai
+Description-Content-Type: text/markdown
+
+# vectorize-openai
+
+Simple wrapper of OpenAI for vectorize requests with single request.
+
+## Installation
+
+```bash
+pip install git+https://github.com/anaregdesign/vectorize-openai.git
+```
+
+## Uninstall
+
+```bash
+pip uninstall openaivec
+```
+
+## Basic Usage
+
+```python
+import os
+from openai import AzureOpenAI
+from openaivec import VectorizedOpenAI
+
+os.environ["AZURE_OPENAI_API_KEY"] = "<your_api_key>"
+api_version = "2024-10-21"
+azure_endpoint = "https://<your_resource_name>.openai.azure.com"
+deployment_name = "<your_deployment_name>"
+
+client = VectorizedOpenAI(
+    client=AzureOpenAI(
+        api_version=api_version,
+        azure_endpoint=azure_endpoint
+    ),
+    temperature=0.0,
+    top_p=1.0,
+    model_name=deployment_name,
+    system_message="Please answer simply with a simple “xx family” and do not output anything else."
+)
+
+client.predict(["panda", "rabit", "koala"])  # => ['bear family', 'rabbit family', 'koala family']
+```
+
+## Usage, process with pandas
+
+```python
+import pandas as pd
+
+...
+
+df = pd.DataFrame({"name": ["panda", "rabbit", "koala"]})
+
+df.assign(
+    kind=lambda df: client.predict(df.name)
+)
+```
+
+the result is:
+
+| name   | kind          |
+|--------|---------------|
+| panda  | bear family   |
+| rabbit | rabbit family |
+| koala  | koala family  |
+
+## Using Azure OpenAI with Apache Spark UDF
+
+Here's simple example of parsing product names using OpenAI with Apache Spark UDF.
+
+You can use the `openaivec` package to create a UDF function to use with Apache Spark.
+At first, you need to create a `UDFConfig` object with the configuration of your OpenAI deployment.
+
+```python
+from openaivec.spark import UDFBuilder
+
+udf = UDFBuilder(
+    api_key="<your-api-key>",
+    api_version="2024-10-21",
+    endpoint="https://<your-resource-name>.openai.azure.com",
+    model_name="<your-deployment-name"
+)
+
+```
+
+here you can use the `completion_udf` function to create a UDF function to use with Apache Spark.
+
+```python
+spark.udf.register("parse_taste", udf.completion("""
+- Extract flavor-related information included in the product name. Only output the flavor name concisely, and nothing else.  
+- Minimize unnecessary adjectives regarding the flavor as much as possible.  
+    - Example:  
+        - Hokkaido Milk → Milk  
+        - Uji Matcha → Matcha  
+
+"""))
+
+spark.udf.register("parse_product", udf.completion("""
+- Extract the type of food included in the product name. Only output the food category and nothing else.  
+- Example output:  
+    - Smoothie  
+    - Milk Tea  
+    - Protein Bar  
+"""))
+```
+
+and then you can use the UDF function in your queries.
+
+```sparksql
+select id,
+       product_name,
+       parse_taste(product_name)   as taste,
+       parse_product(product_name) as product
+from product_names
+```
+
+Output:
+
+| id            | product_name                         | taste     | product     |
+|---------------|--------------------------------------|-----------|-------------|
+| 4414732714624 | Cafe Mocha Smoothie (Trial Size)     | Mocha     | Smoothie    |
+| 4200162318339 | Dark Chocolate Tea (New Product)     | Chocolate | Tea         |
+| 4920122084098 | Cafe Mocha Protein Bar (Trial Size)  | Mocha     | Protein Bar |
+| 4468864478874 | Dark Chocolate Smoothie (On Sale)    | Chocolate | Smoothie    |
+| 4036242144725 | Uji Matcha Tea (New Product)         | Matcha    | Tea         |
+| 4847798245741 | Hokkaido Milk Tea (Trial Size)       | Milk      | Milk Tea    |
+| 4449574211957 | Dark Chocolate Smoothie (Trial Size) | Chocolate | Smoothie    |
+| 4127044426148 | Fruit Mix Tea (Trial Size)           | Fruit     | Tea         |
+| ...           | ...                                  | ...       | ...         |
+
+
+
+
+

openaivec-0.2.1/README.md

@@ -0,0 +1,131 @@
+# vectorize-openai
+
+Simple wrapper of OpenAI for vectorize requests with single request.
+
+## Installation
+
+```bash
+pip install git+https://github.com/anaregdesign/vectorize-openai.git
+```
+
+## Uninstall
+
+```bash
+pip uninstall openaivec
+```
+
+## Basic Usage
+
+```python
+import os
+from openai import AzureOpenAI
+from openaivec import VectorizedOpenAI
+
+os.environ["AZURE_OPENAI_API_KEY"] = "<your_api_key>"
+api_version = "2024-10-21"
+azure_endpoint = "https://<your_resource_name>.openai.azure.com"
+deployment_name = "<your_deployment_name>"
+
+client = VectorizedOpenAI(
+    client=AzureOpenAI(
+        api_version=api_version,
+        azure_endpoint=azure_endpoint
+    ),
+    temperature=0.0,
+    top_p=1.0,
+    model_name=deployment_name,
+    system_message="Please answer simply with a simple “xx family” and do not output anything else."
+)
+
+client.predict(["panda", "rabit", "koala"])  # => ['bear family', 'rabbit family', 'koala family']
+```
+
+## Usage, process with pandas
+
+```python
+import pandas as pd
+
+...
+
+df = pd.DataFrame({"name": ["panda", "rabbit", "koala"]})
+
+df.assign(
+    kind=lambda df: client.predict(df.name)
+)
+```
+
+the result is:
+
+| name   | kind          |
+|--------|---------------|
+| panda  | bear family   |
+| rabbit | rabbit family |
+| koala  | koala family  |
+
+## Using Azure OpenAI with Apache Spark UDF
+
+Here's simple example of parsing product names using OpenAI with Apache Spark UDF.
+
+You can use the `openaivec` package to create a UDF function to use with Apache Spark.
+At first, you need to create a `UDFConfig` object with the configuration of your OpenAI deployment.
+
+```python
+from openaivec.spark import UDFBuilder
+
+udf = UDFBuilder(
+    api_key="<your-api-key>",
+    api_version="2024-10-21",
+    endpoint="https://<your-resource-name>.openai.azure.com",
+    model_name="<your-deployment-name"
+)
+
+```
+
+here you can use the `completion_udf` function to create a UDF function to use with Apache Spark.
+
+```python
+spark.udf.register("parse_taste", udf.completion("""
+- Extract flavor-related information included in the product name. Only output the flavor name concisely, and nothing else.  
+- Minimize unnecessary adjectives regarding the flavor as much as possible.  
+    - Example:  
+        - Hokkaido Milk → Milk  
+        - Uji Matcha → Matcha  
+
+"""))
+
+spark.udf.register("parse_product", udf.completion("""
+- Extract the type of food included in the product name. Only output the food category and nothing else.  
+- Example output:  
+    - Smoothie  
+    - Milk Tea  
+    - Protein Bar  
+"""))
+```
+
+and then you can use the UDF function in your queries.
+
+```sparksql
+select id,
+       product_name,
+       parse_taste(product_name)   as taste,
+       parse_product(product_name) as product
+from product_names
+```
+
+Output:
+
+| id            | product_name                         | taste     | product     |
+|---------------|--------------------------------------|-----------|-------------|
+| 4414732714624 | Cafe Mocha Smoothie (Trial Size)     | Mocha     | Smoothie    |
+| 4200162318339 | Dark Chocolate Tea (New Product)     | Chocolate | Tea         |
+| 4920122084098 | Cafe Mocha Protein Bar (Trial Size)  | Mocha     | Protein Bar |
+| 4468864478874 | Dark Chocolate Smoothie (On Sale)    | Chocolate | Smoothie    |
+| 4036242144725 | Uji Matcha Tea (New Product)         | Matcha    | Tea         |
+| 4847798245741 | Hokkaido Milk Tea (Trial Size)       | Milk      | Milk Tea    |
+| 4449574211957 | Dark Chocolate Smoothie (Trial Size) | Chocolate | Smoothie    |
+| 4127044426148 | Fruit Mix Tea (Trial Size)           | Fruit     | Tea         |
+| ...           | ...                                  | ...       | ...         |
+
+
+
+

openaivec-0.2.1/openaivec/__init__.py

@@ -0,0 +1,7 @@
+from .embedding import EmbeddingOpenAI
+from .vectorize import VectorizedOpenAI
+
+__ALL__ = [
+    "VectorizedOpenAI",
+    "EmbeddingOpenAI",
+]

openaivec-0.2.1/openaivec/embedding.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+from typing import List
+
+import numpy as np
+from numpy.typing import NDArray
+from openai import OpenAI
+
+from openaivec.util import map_unique_minibatch
+
+__ALL__ = ["EmbeddingOpenAI"]
+
+
+@dataclass(frozen=True)
+class EmbeddingOpenAI:
+    client: OpenAI
+    model_name: str
+
+    def embed(self, sentences: List[str]) -> List[NDArray[np.float32]]:
+        responses = self.client.embeddings.create(input=sentences, model=self.model_name)
+        return [np.array(d.embedding, dtype=np.float32) for d in responses.data]
+
+    def embed_minibatch(self, sentences: List[str], batch_size: int) -> List[NDArray[np.float32]]:
+        return map_unique_minibatch(sentences, batch_size, self.embed)

openaivec-0.2.1/openaivec/spark.py

@@ -0,0 +1,91 @@
+import os
+from dataclasses import dataclass
+from typing import Iterator
+
+import pandas as pd
+from pyspark.sql.pandas.functions import pandas_udf
+from pyspark.sql.types import StringType, ArrayType, FloatType
+
+__ALL__ = ["UDFBuilder"]
+
+
+@dataclass(frozen=True)
+class UDFBuilder:
+    api_key: str
+    api_version: str
+    endpoint: str
+    model_name: str
+    batch_size: int = 256
+
+    @classmethod
+    def of_environment(cls) -> "UDFBuilder":
+        return cls(
+            api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
+            api_version=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-10-21"),
+            endpoint=os.environ.get("AZURE_OPENAI_ENDPOINT"),
+            model_name=os.environ.get("AZURE_OPENAI_MODEL_NAME"),
+        )
+
+    def __post_init__(self):
+        assert self.api_key, "api_key must be set"
+        assert self.api_version, "api_version must be set"
+        assert self.endpoint, "endpoint must be set"
+        assert self.model_name, "model_name must be set"
+
+    def completion(self, system_message: str):
+        @pandas_udf(StringType())
+        def fn(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
+            import httpx
+            import pandas as pd
+            from openai import AzureOpenAI
+
+            from openaivec import VectorizedOpenAI
+
+            client = AzureOpenAI(
+                api_version=self.api_version,
+                azure_endpoint=self.endpoint,
+                http_client=httpx.Client(http2=True, verify=False),
+                api_key=self.api_key,
+            )
+
+            client_vec = VectorizedOpenAI(
+                client=client,
+                model_name=self.model_name,
+                system_message=system_message,
+                top_p=1.0,
+                temperature=0.0,
+            )
+
+            for part in col:
+                yield pd.Series(
+                    client_vec.predict_minibatch(part.tolist(), self.batch_size)
+                )
+
+        return fn
+
+    def embedding(self):
+        @pandas_udf(ArrayType(FloatType()))
+        def fn(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
+            import httpx
+            from openai import AzureOpenAI
+
+            from openaivec.embedding import EmbeddingOpenAI
+
+            client = AzureOpenAI(
+                api_version=self.api_version,
+                azure_endpoint=self.endpoint,
+                http_client=httpx.Client(http2=True, verify=False),
+                api_key=self.api_key,
+            )
+
+            client_emb = EmbeddingOpenAI(
+                client=client,
+                model_name=self.model_name,
+            )
+
+            for part in col:
+                yield pd.Series(
+                    client_emb.embed_minibatch(part.tolist(), self.batch_size)
+                )
+
+        return fn

openaivec-0.2.1/openaivec/test_util.py

@@ -0,0 +1,76 @@
+from typing import List
+from unittest import TestCase
+
+from openaivec.util import split_to_minibatch, map_minibatch, map_unique, map_unique_minibatch
+
+
+class TestMappingFunctions(TestCase):
+
+    def test_split_to_minibatch_normal(self):
+        b = [1, 2, 3, 4, 5]
+        batch_size = 2
+        expected = [[1, 2], [3, 4], [5]]
+        self.assertEqual(split_to_minibatch(b, batch_size), expected)
+
+    def test_split_to_minibatch_empty(self):
+        b: List[int] = []
+        batch_size = 3
+        expected: List[List[int]] = []
+        self.assertEqual(split_to_minibatch(b, batch_size), expected)
+
+    def test_map_minibatch(self):
+        # Function that doubles each element in the batch.
+        def double_list(lst: List[int]) -> List[int]:
+            return [x * 2 for x in lst]
+
+        b = [1, 2, 3, 4, 5]
+        batch_size = 2
+        # Batches: [1,2] -> [2,4], [3,4] -> [6,8], [5] -> [10]
+        expected = [2, 4, 6, 8, 10]
+        self.assertEqual(map_minibatch(b, batch_size, double_list), expected)
+
+    def test_map_minibatch_batch_size_one(self):
+        # Identity function: returns the list as is.
+        def identity(lst: List[int]) -> List[int]:
+            return lst
+
+        b = [1, 2, 3, 4]
+        batch_size = 1
+        expected = [1, 2, 3, 4]
+        self.assertEqual(map_minibatch(b, batch_size, identity), expected)
+
+    def test_map_minibatch_batch_size_greater_than_list(self):
+        def identity(lst: List[int]) -> List[int]:
+            return lst
+
+        b = [1, 2, 3]
+        batch_size = 5
+        expected = [1, 2, 3]
+        self.assertEqual(map_minibatch(b, batch_size, identity), expected)
+
+    def test_map_unique(self):
+        # Function that squares each element.
+        def square_list(lst: List[int]) -> List[int]:
+            return [x * x for x in lst]
+
+        b = [3, 2, 3, 1]
+        # Unique order preserved using dict.fromkeys: [3, 2, 1]
+        # After applying f: [9, 4, 1]
+        # Mapping back for original list: [9, 4, 9, 1]
+        expected = [9, 4, 9, 1]
+        self.assertEqual(map_unique(b, square_list), expected)
+
+    def test_map_unique_minibatch(self):
+        # Function that doubles each element.
+        def double_list(lst: List[int]) -> List[int]:
+            return [x * 2 for x in lst]
+
+        b = [1, 2, 1, 3]
+        batch_size = 2
+        # Unique order: [1, 2, 3]
+        # Using map_minibatch on unique values:
+        #  Split [1,2,3] with batch_size=2 -> [[1,2], [3]]
+        #  Apply function: [[2,4], [6]] -> flattened to [2,4,6]
+        # Mapping back for original list: [2, 4, 2, 6]
+        expected = [2, 4, 2, 6]
+        self.assertEqual(map_unique_minibatch(b, batch_size, double_list), expected)

openaivec-0.2.1/openaivec/util.py

@@ -0,0 +1,38 @@
+from itertools import chain
+from typing import List, TypeVar, Callable
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+def split_to_minibatch(b: List[T], batch_size: int) -> List[List[T]]:
+    """Splits the list into sublists of size `batch_size`."""
+    return [b[i:i + batch_size] for i in range(0, len(b), batch_size)]
+
+def map_minibatch(b: List[T], batch_size: int, f: Callable[[List[T]], List[U]]) -> List[U]:
+    """
+    Splits the list `b` into batches of size `batch_size` and applies the function `f` to each batch.
+    The results (each a list) are then flattened into a single list.
+    """
+    batches = split_to_minibatch(b, batch_size)
+    return list(chain.from_iterable(f(batch) for batch in batches))
+
+def map_unique(b: List[T], f: Callable[[List[T]], List[U]]) -> List[U]:
+    """
+    Applies the function `f` only once to the unique values in the list `b` (preserving their order),
+    and then maps the resulting values back to match the original list.
+    This avoids repeated execution of `f` for duplicate values.
+    """
+    # Use dict.fromkeys to remove duplicates while preserving the order
+    unique_values = list(dict.fromkeys(b))
+    value_to_index = {v: i for i, v in enumerate(unique_values)}
+    results = f(unique_values)
+    return [results[value_to_index[value]] for value in b]
+
+
+def map_unique_minibatch(b: List[T], batch_size: int, f: Callable[[List[T]], List[U]]) -> List[U]:
+    """
+    Uses minibatch processing on the unique values of the list `b`.
+    The function `f` is applied to these unique values in batches,
+    and the results are mapped back to match the order of the original list.
+    """
+    return map_unique(b, lambda x: map_minibatch(x, batch_size, f))

openaivec-0.2.1/openaivec/vectorize.py

@@ -0,0 +1,108 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from openai import OpenAI
+from openai.types.chat import ParsedChatCompletion
+from pydantic import BaseModel
+
+from openaivec.util import map_unique_minibatch
+
+__ALL__ = ["VectorizedOpenAI"]
+
+
+def vectorize_system_message(system_message: str) -> str:
+    return f"""
+<SystemMessage>
+    <Instructions>
+        <Instruction>{system_message}</Instruction>
+        <Instruction>
+            You will receive multiple user messages at once.
+            Please provide an appropriate response to each message individually.
+        </Instruction>
+    </Instructions>
+    <Examples>
+        <Example>
+            <Input>
+                {{
+                    "user_messages": [
+                        {{
+                            "id": 1,
+                            "text": "{{user_message_1}}"
+                        }},
+                        {{
+                            "id": 2,
+                            "text": "{{user_message_2}}"
+                        }}
+                    ]
+                }}
+            </Input>
+            <Output>
+                {{
+                    "assistant_messages": [
+                        {{
+                            "id": 1,
+                            "text": "{{assistant_response_1}}"
+                        }},
+                        {{
+                            "id": 2,
+                            "text": "{{assistant_response_2}}"
+                        }}
+                    ]
+                }}
+            </Output>
+        </Example>
+    </Examples>
+</SystemMessage>
+"""
+
+
+class Message(BaseModel):
+    id: int
+    text: str
+
+
+class Request(BaseModel):
+    user_messages: List[Message]
+
+
+class Response(BaseModel):
+    assistant_messages: List[Message]
+
+
+@dataclass(frozen=True)
+class VectorizedOpenAI:
+    client: OpenAI
+    model_name: str  ## it would be the name of deployment for Azure
+    system_message: str
+    temperature: float = 0.0
+    top_p: float = 1.0
+    _vectorized_system_message: str = field(init=False)
+
+    def __post_init__(self):
+        object.__setattr__(self, "_vectorized_system_message", vectorize_system_message(self.system_message))
+
+    def request(self, user_messages: List[Message]) -> ParsedChatCompletion[Response]:
+        completion = self.client.beta.chat.completions.parse(
+            model=self.model_name,
+            messages=[
+                {"role": "system", "content": self._vectorized_system_message},
+                {"role": "user", "content": Request(user_messages=user_messages).model_dump_json()}
+            ],
+            temperature=self.temperature,
+            top_p=self.top_p,
+            response_format=Response
+        )
+        return completion
+
+    def predict(self, user_messages: List[str]) -> List[str]:
+        messages = [Message(id=i, text=message) for i, message in enumerate(user_messages)]
+        completion = self.request(messages)
+        response_dict = {
+            message.id: message.text
+            for message in completion.choices[0].message.parsed.assistant_messages
+        }
+        sorted_responses = [response_dict[m.id] for m in messages]
+        return sorted_responses
+
+    def predict_minibatch(self, user_messages: List[str], batch_size: int) -> List[str]:
+        return map_unique_minibatch(user_messages, batch_size, self.predict)

openaivec-0.2.1/pyproject.toml

@@ -0,0 +1,27 @@
+[tool.poetry]
+name = "openaivec"
+version = "0.2.1"
+description = ""
+authors = ["Hiroki Mizukami <hmizukami@microsoft.com>"]
+license = "MIT License"
+readme = "README.md"
+homepage = "https://github.com/anaregdesign/vectorize-openai"
+repository = "https://github.com/anaregdesign/vectorize-openai"
+packages = [
+    { include = "openaivec", from = "." }
+]
+
+[tool.poetry.dependencies]
+python = "^3.10"
+pandas = "^2.2.3"
+pyspark = "^3.5.1"
+openai = "^1.57.2"
+httpx = {extras = ["http2"], version = "^0.28.1"}
+
+[tool.poetry.dev-dependencies]
+pytest = "^8.3.4"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"