# This file was auto-generated by Fern from our API Definition.

from ..core.pydantic_utilities import UniversalBaseModel
import typing
import pydantic
from .pooling import Pooling
from .token import Token
from ..core.pydantic_utilities import IS_PYDANTIC_V2


class HuggingFaceInferenceApiEmbedding(UniversalBaseModel):
    model_name: typing.Optional[str] = pydantic.Field(default=None)
    """
    Hugging Face model name. If None, the task will be used.
    """

    embed_batch_size: typing.Optional[int] = pydantic.Field(default=None)
    """
    The batch size for embedding calls.
    """

    num_workers: typing.Optional[int] = pydantic.Field(default=None)
    """
    The number of workers to use for async embedding calls.
    """

    pooling: typing.Optional[Pooling] = pydantic.Field(default=None)
    """
    Pooling strategy. If None, the model's default pooling is used.
    """

    query_instruction: typing.Optional[str] = pydantic.Field(default=None)
    """
    Instruction to prepend during query embedding.
    """

    text_instruction: typing.Optional[str] = pydantic.Field(default=None)
    """
    Instruction to prepend during text embedding.
    """

    token: typing.Optional[Token] = pydantic.Field(default=None)
    """
    Hugging Face token. Will default to the locally saved token. Pass token=False if you don’t want to send your token to the server.
    """

    timeout: typing.Optional[float] = pydantic.Field(default=None)
    """
    The maximum number of seconds to wait for a response from the server. Loading a new model in Inference API can take up to several minutes. Defaults to None, meaning it will loop until the server is available.
    """

    headers: typing.Optional[typing.Dict[str, typing.Optional[str]]] = pydantic.Field(default=None)
    """
    Additional headers to send to the server. By default only the authorization and user-agent headers are sent. Values in this dictionary will override the default values.
    """

    cookies: typing.Optional[typing.Dict[str, typing.Optional[str]]] = pydantic.Field(default=None)
    """
    Additional cookies to send to the server.
    """

    task: typing.Optional[str] = pydantic.Field(default=None)
    """
    Optional task to pick Hugging Face's recommended model, used when model_name is left as default of None.
    """

    class_name: typing.Optional[str] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
    else:

        class Config:
            frozen = True
            smart_union = True
            extra = pydantic.Extra.allow
