diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 87f9a7a..6ff4011 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -49,6 +49,8 @@ jobs: - name: Run tests run: | poetry run pytest -v -s ./tests + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} lint_check: runs-on: ubuntu-latest diff --git a/tests/embedders/test_openai.py b/tests/embedders/test_openai.py new file mode 100644 index 0000000..f33eb4d --- /dev/null +++ b/tests/embedders/test_openai.py @@ -0,0 +1,23 @@ +import numpy as np + +from jmteb.embedders.openai_embedder import OpenAIEmbedder + +OUTPUT_DIM = 1536 + + +class TestSentenceBertEmbedder: + def setup_class(cls): + cls.model = OpenAIEmbedder(model="text-embedding-3-small") + + def test_encode(self): + embeddings = self.model.encode("任意のテキスト") + assert isinstance(embeddings, np.ndarray) + assert embeddings.shape == (OUTPUT_DIM,) + + def test_encode_multiple(self): + embeddings = self.model.encode(["任意のテキスト"] * 3) + assert isinstance(embeddings, np.ndarray) + assert embeddings.shape == (3, OUTPUT_DIM) + + def test_get_output_dim(self): + assert self.model.get_output_dim() == OUTPUT_DIM