Merge pull request #2 from dataforgoodfr/feature/GenAI-client-with-du…

…mmy-impacts-values

Energy estimation for OpenAI

genai_impact/__init__.py

@@ -0,0 +1 @@
+from .client_wrapper import OpenAI

genai_impact/__main__.py

@@ -0,0 +1,12 @@
+from genai_impact.client_wrapper import OpenAI
+
+client = OpenAI()
+
+response = client.chat.completions.create(
+    model="gpt-3.5-turbo",
+    messages=[
+        {"role": "user", "content": "Hello World!"}
+    ]
+)
+
+print(response.impacts)

genai_impact/client_wrapper.py

@@ -0,0 +1,43 @@
+from dataclasses import asdict
+from typing import Optional
+
+from openai import OpenAI as _OpenAI
+from wrapt import wrap_function_wrapper
+
+from genai_impact.compute_impacts import MODEL_SIZES, compute_llm_impact
+
+
+def chat_wrapper(wrapped, instance, args, kwargs):
+    response = wrapped(*args, **kwargs)
+    model_size = MODEL_SIZES.get(response.model)
+    output_tokens = response.usage.completion_tokens
+    impacts = compute_llm_impact(model_parameter_count=model_size, output_token_count=output_tokens)
+    response.impacts = asdict(impacts)
+    return response
+
+
+class OpenAI:
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        **kwargs
+    ):
+        self.__client = _OpenAI(
+            api_key=api_key,
+            base_url=base_url,
+            **kwargs
+        )
+
+    wrap_function_wrapper(
+        "openai.resources.chat.completions",
+        "Completions.create",
+        chat_wrapper
+    )
+
+    def __getattr__(self, name):
+        """
+        Redirect attribute access to the underlying openai client if the attribute
+        is not defined in this class.
+        """
+        return getattr(self.__client, name)

genai_impact/compute_impacts.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+
+
+ENERGY_PROFILE = 1.17e-4
+
+MODEL_SIZES = {
+    "gpt-4-0125-preview": None,
+    "gpt-4-turbo-preview": None,
+    "gpt-4-1106-preview": None,
+    "gpt-4-vision-preview": None,
+    "gpt-4": 200,
+    "gpt-4-0314": 200,
+    "gpt-4-0613": 200,
+    "gpt-4-32k": 200,
+    "gpt-4-32k-0314": 200,
+    "gpt-4-32k-0613": 200,
+    "gpt-3.5-turbo": 20,
+    "gpt-3.5-turbo-16k": 20,
+    "gpt-3.5-turbo-0301": 20,
+    "gpt-3.5-turbo-0613": 20,
+    "gpt-3.5-turbo-1106": 20,
+    "gpt-3.5-turbo-0125": 20,
+    "gpt-3.5-turbo-16k-0613": 20,
+}
+
+
+@dataclass
+class Impacts:
+    energy: float
+    energy_unit: str = 'Wh'
+
+
+def compute_llm_impact(
+    model_parameter_count: float,
+    output_token_count: int,
+) -> Impacts:
+    return Impacts(
+        energy=ENERGY_PROFILE * model_parameter_count * output_token_count
+    )

pyproject.toml

@@ -12,6 +12,8 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = ">=3.8,<4"
+openai = "^1.12.0"
+wrapt = "^1.16.0"
 
 
 [tool.poetry.group.dev.dependencies]

tests/test_compute_impacts.py

@@ -0,0 +1,12 @@
+import pytest
+
+from genai_impact.compute_impacts import compute_llm_impact
+
+
+@pytest.mark.parametrize('model_size,output_tokens', [
+    (130, 1000),
+    (7, 150)
+])
+def test_compute_impacts_is_positive(model_size: float, output_tokens: int) -> None:
+    impacts = compute_llm_impact(model_size, output_tokens)
+    assert impacts.energy >= 0