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

import datetime as dt
import typing

from ..core.datetime_utils import serialize_datetime
from .hugging_face_inference_api_embedding_token import HuggingFaceInferenceApiEmbeddingToken
from .pooling import Pooling

try:
    import pydantic.v1 as pydantic  # type: ignore
except ImportError:
    import pydantic  # type: ignore


class HuggingFaceInferenceApiEmbedding(pydantic.BaseModel):
    """
    Wrapper on the Hugging Face's Inference API for embeddings.

    Overview of the design:

    - Uses the feature extraction task: https://huggingface.co/tasks/feature-extraction
    """

    model_name: typing.Optional[str] = pydantic.Field(
        description="The model to run inference with. Can be a model id hosted on the Hugging Face Hub, e.g. bigcode/starcoder or a URL to a deployed Inference Endpoint. Defaults to None, in which case a recommended model is automatically selected for the task (see Field below)."
    )
    embed_batch_size: typing.Optional[int] = pydantic.Field(description="The batch size for embedding calls.")
    callback_manager: typing.Optional[typing.Dict[str, typing.Any]]
    pooling: typing.Optional[Pooling] = pydantic.Field(
        description="Optional pooling technique to use with embeddings capability, if the model's raw output needs pooling."
    )
    query_instruction: typing.Optional[str] = pydantic.Field(
        description="Instruction to prepend during query embedding. Use of None means infer the instruction based on the model. Use of empty string will defeat instruction prepending entirely."
    )
    text_instruction: typing.Optional[str] = pydantic.Field(
        description="Instruction to prepend during text embedding. Use of None means infer the instruction based on the model. Use of empty string will defeat instruction prepending entirely."
    )
    token: typing.Optional[HuggingFaceInferenceApiEmbeddingToken] = pydantic.Field(
        description="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(
        description="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, str]] = pydantic.Field(
        description="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, str]] = pydantic.Field(
        description="Additional cookies to send to the server."
    )
    task: typing.Optional[str] = pydantic.Field(
        description="Optional task to pick Hugging Face's recommended model, used when model_name is left as default of None."
    )
    class_name: typing.Optional[str]

    def json(self, **kwargs: typing.Any) -> str:
        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
        return super().json(**kwargs_with_defaults)

    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
        return super().dict(**kwargs_with_defaults)

    class Config:
        frozen = True
        smart_union = True
        json_encoders = {dt.datetime: serialize_datetime}
