File size: 2,225 Bytes
e3278e4 |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
"""
Nvidia NIM embeddings endpoint: https://docs.api.nvidia.com/nim/reference/nvidia-nv-embedqa-e5-v5-infer
This is OpenAI compatible
This file only contains param mapping logic
API calling is done using the OpenAI SDK with an api_base
"""
import types
from typing import Optional
class NvidiaNimEmbeddingConfig:
"""
Reference: https://docs.api.nvidia.com/nim/reference/nvidia-nv-embedqa-e5-v5-infer
"""
# OpenAI params
encoding_format: Optional[str] = None
user: Optional[str] = None
# Nvidia NIM params
input_type: Optional[str] = None
truncate: Optional[str] = None
def __init__(
self,
encoding_format: Optional[str] = None,
user: Optional[str] = None,
input_type: Optional[str] = None,
truncate: Optional[str] = None,
) -> None:
locals_ = locals()
for key, value in locals_.items():
if key != "self" and value is not None:
setattr(self.__class__, key, value)
@classmethod
def get_config(cls):
return {
k: v
for k, v in cls.__dict__.items()
if not k.startswith("__")
and not isinstance(
v,
(
types.FunctionType,
types.BuiltinFunctionType,
classmethod,
staticmethod,
),
)
and v is not None
}
def get_supported_openai_params(
self,
):
return ["encoding_format", "user"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
kwargs: Optional[dict] = None,
):
if "extra_body" not in optional_params:
optional_params["extra_body"] = {}
for k, v in non_default_params.items():
if k == "input_type":
optional_params["extra_body"].update({"input_type": v})
elif k == "truncate":
optional_params["extra_body"].update({"truncate": v})
if kwargs is not None:
# pass kwargs in extra_body
optional_params["extra_body"].update(kwargs)
return optional_params
|