File size: 1,623 Bytes
5d3967e 11fa46b 5d3967e 11fa46b 5d3967e 11fa46b 5d3967e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import os
from typing import Literal
import pytest
from openai import OpenAI
from openai.types import CreateEmbeddingResponse
@pytest.fixture
def client():
return OpenAI(
base_url=os.environ.get("HFENDPOINTS_BASE_URL", "http://localhost:8000/api/v1/"),
api_key=os.environ.get("HFENDPOINTS_API_KEY", os.environ.get("HF_TOKEN", "hf_notdefined")),
organization="hfendpoints-images",
)
@pytest.mark.parametrize("encoding_format", ["float", "base64"])
def test_openai_embedding_single(client, encoding_format: Literal["float", "base64"]):
response: CreateEmbeddingResponse = client.embeddings.create(
input="Hello, how are you?",
model="all-MiniLM-L6-v2",
encoding_format=encoding_format,
)
assert len(response.model)
assert response.object == "list"
assert len(response.data) == 1
assert response.data[0].object == "embedding"
assert len(response.data[0].embedding) == 384
@pytest.mark.parametrize("encoding_format", ["float", "base64"])
def test_openai_embedding_batch(client, encoding_format: Literal["float", "base64"]):
response: CreateEmbeddingResponse = client.embeddings.create(
input=["Hello, how are you?", "How are you?"],
model="all-MiniLM-L6-v2",
encoding_format=encoding_format,
)
assert len(response.model)
assert response.object == "list"
assert len(response.data) == 2
assert response.data[0].object == "embedding"
assert len(response.data[0].embedding) == 384
assert response.data[1].object == "embedding"
assert len(response.data[1].embedding) == 384 |