98 lines
2.4 KiB
Python
98 lines
2.4 KiB
Python
from typing import Optional, Union
|
|
|
|
from mlflow.gateway.base_models import RequestModel, ResponseModel
|
|
from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
|
|
|
|
_REQUEST_PAYLOAD_EXTRA_SCHEMA = {
|
|
"example": {
|
|
"input": ["hello", "world"],
|
|
}
|
|
}
|
|
|
|
|
|
class RequestPayload(RequestModel):
|
|
input: Union[str, list[str], list[int], list[list[int]]]
|
|
|
|
class Config:
|
|
if IS_PYDANTIC_V2_OR_NEWER:
|
|
json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
|
|
else:
|
|
schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
|
|
|
|
|
|
class EmbeddingObject(ResponseModel):
|
|
object: str = "embedding"
|
|
embedding: Union[list[float], str]
|
|
index: int
|
|
|
|
|
|
class EmbeddingsUsage(ResponseModel):
|
|
prompt_tokens: Optional[int] = None
|
|
total_tokens: Optional[int] = None
|
|
|
|
|
|
_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"object": "embedding",
|
|
"index": 0,
|
|
"embedding": [
|
|
0.017291732,
|
|
-0.017291732,
|
|
0.014577783,
|
|
-0.02902633,
|
|
-0.037271563,
|
|
0.019333655,
|
|
-0.023055641,
|
|
-0.007359971,
|
|
-0.015818445,
|
|
-0.030654699,
|
|
0.008348623,
|
|
0.018312693,
|
|
-0.017149571,
|
|
-0.0044424757,
|
|
-0.011165961,
|
|
0.01018377,
|
|
],
|
|
},
|
|
{
|
|
"object": "embedding",
|
|
"index": 1,
|
|
"embedding": [
|
|
0.0060126893,
|
|
-0.008691099,
|
|
-0.0040095365,
|
|
0.019889368,
|
|
0.036211833,
|
|
-0.0013270887,
|
|
0.013401738,
|
|
-0.0036735237,
|
|
-0.0049594184,
|
|
0.035229642,
|
|
-0.03435084,
|
|
0.019798903,
|
|
-0.0006110424,
|
|
0.0073793563,
|
|
0.005657291,
|
|
0.022487005,
|
|
],
|
|
},
|
|
],
|
|
"model": "text-embedding-ada-002-v2",
|
|
"usage": {"prompt_tokens": 400, "total_tokens": 400},
|
|
}
|
|
|
|
|
|
class ResponsePayload(ResponseModel):
|
|
object: str = "list"
|
|
data: list[EmbeddingObject]
|
|
model: str
|
|
usage: EmbeddingsUsage
|
|
|
|
class Config:
|
|
if IS_PYDANTIC_V2_OR_NEWER:
|
|
json_schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
|
|
else:
|
|
schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
|