From 356fcb72f6def4067e041b64d25678718da42c2a Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Thu, 14 Sep 2023 14:01:27 -0700
Subject: [PATCH] NumPy decoding docs, plus extra tests for llm.encode/decode

!stable-docs

Refs https://discord.com/channels/823971286308356157/1128504153841336370/1151975583237034056
---
 docs/embeddings/storage.md  |  9 +++++++++
 setup.py                    |  1 +
 tests/test_encode_decode.py | 19 +++++++++++++++++++
 3 files changed, 29 insertions(+)
 create mode 100644 tests/test_encode_decode.py

diff --git a/docs/embeddings/storage.md b/docs/embeddings/storage.md
index da99cdc..f932b47 100644
--- a/docs/embeddings/storage.md
+++ b/docs/embeddings/storage.md
@@ -20,3 +20,12 @@ def decode(binary):
 ```
 
 These functions are available as `llm.encode()` and `llm.decode()`.
+
+If you are using [NumPy](https://numpy.org/) you can decode one of these binary values like this:
+
+```python
+import numpy as np
+
+numpy_array = np.frombuffer(value, "<f4")
+```
+The `<f4` format string here ensures NumPy will treat the data as a little-endian sequence of 32-bit floats.
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 8025e68..6dbbd7f 100644
--- a/setup.py
+++ b/setup.py
@@ -51,6 +51,7 @@ setup(
     extras_require={
         "test": [
             "pytest",
+            "numpy",
             "requests-mock",
             "cogapp",
             "mypy",
diff --git a/tests/test_encode_decode.py b/tests/test_encode_decode.py
new file mode 100644
index 0000000..1310c6d
--- /dev/null
+++ b/tests/test_encode_decode.py
@@ -0,0 +1,19 @@
+import llm
+import pytest
+import numpy as np
+
+
+@pytest.mark.parametrize(
+    "array",
+    (
+        (0.0, 1.0, 1.5),
+        (3423.0, 222.0, -1234.5),
+    ),
+)
+def test_roundtrip(array):
+    encoded = llm.encode(array)
+    decoded = llm.decode(encoded)
+    assert decoded == array
+    # Try with numpy as well
+    numpy_decoded = np.frombuffer(encoded, "<f4")
+    assert tuple(numpy_decoded.tolist()) == array