# type: ignore
# Generated by the protocol buffer compiler.  DO NOT EDIT!
# sources: arg_services/nlp/v1/nlp.proto
# plugin: python-betterproto
# This file has been @generated

from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Dict,
    List,
    Optional,
)

import betterproto
import betterproto.lib.google.protobuf as betterproto_lib_google_protobuf
import grpclib
from betterproto.grpc.grpclib_server import ServiceBase


if TYPE_CHECKING:
    import grpclib.server
    from betterproto.grpc.grpclib_client import MetadataLike
    from grpclib.metadata import Deadline


class SimilarityMethod(betterproto.Enum):
    """Possible methods to compute the similarity between two vectors."""

    UNSPECIFIED = 0
    """If not given, the implementation defaults to cosine similarity."""

    COSINE = 1
    """
    Cosine similarity.
    [Wikipedia](https://en.wikipedia.org/wiki/Cosine_similarity).
    """

    DYNAMAX_JACCARD = 2
    """
    DynaMax Jaccard. [Paper](https://arxiv.org/abs/1904.13264), [Code](https://
    github.com/babylonhealth/fuzzymax/blob/master/similarity/fuzzy.py).
    """

    MAXPOOL_JACCARD = 3
    """
    MaxPool Jaccard. [Paper](https://arxiv.org/abs/1904.13264), [Code](https://
    github.com/babylonhealth/fuzzymax/blob/master/similarity/fuzzy.py).
    """

    DYNAMAX_DICE = 4
    """
    DynaMax Dice. [Paper](https://arxiv.org/abs/1904.13264), [Code](https://git
    hub.com/babylonhealth/fuzzymax/blob/master/similarity/fuzzy.py).
    """

    DYNAMAX_OTSUKA = 5
    """
    DynaMax Otsuka. [Paper](https://arxiv.org/abs/1904.13264), [Code](https://g
    ithub.com/babylonhealth/fuzzymax/blob/master/similarity/fuzzy.py).
    """

    WMD = 6
    """
    Word Mover's Distance [Gensim Tutorial](https://radimrehurek.com/gensim/aut
    o_examples/tutorials/run_wmd.html).
    """

    EDIT = 7
    """
    Levenshtein distance.
    [Wikipedia](https://en.wikipedia.org/wiki/Levenshtein_distance).
    """

    JACCARD = 8
    """
    Jaccard similarity.
    [Wikipedia](https://en.wikipedia.org/wiki/Jaccard_index).
    """

    ANGULAR = 9
    """
    Angular distance.
    [Wikipedia](https://en.wikipedia.org/wiki/Angular_distance).
    """


class EmbeddingLevel(betterproto.Enum):
    UNSPECIFIED = 0
    """In the default case, no vector is computed."""

    DOCUMENT = 1
    """Compute one vector for the whole string."""

    TOKENS = 2
    """Compute vectors for all tokens in the string."""

    SENTENCES = 3
    """Compute vectors for all sentences found in the string."""


class Pooling(betterproto.Enum):
    UNSPECIFIED = 0
    """IN the default case, the arithmetic mean should be used."""

    MEAN = 1
    """
    Arithmetic mean of all elements.
    [Wikipedia](https://en.wikipedia.org/wiki/Arithmetic_mean).
    """

    MAX = 2
    """
    Maximum element of vector.
    [Wikipedia](https://en.wikipedia.org/wiki/Maximum).
    """

    MIN = 3
    """
    Minimum element of vector.
    [Wikipedia](https://en.wikipedia.org/wiki/Minimum).
    """

    SUM = 4
    """Sum of all elements."""

    FIRST = 5
    """First element of vector."""

    LAST = 6
    """Last element of vector."""

    MEDIAN = 7
    """
    Median element of vector.
    [Wikipedia](https://en.wikipedia.org/wiki/Median).
    """

    GMEAN = 8
    """
    Geometirc mean of all elements.
    [Wikipedia](https://en.wikipedia.org/wiki/Geometric_mean).
    """

    HMEAN = 9
    """
    Harmonic mean of all elements.
    [Wikipedia](https://en.wikipedia.org/wiki/Harmonic_mean).
    """


class EmbeddingType(betterproto.Enum):
    UNSPECIFIED = 0
    """In the default case, no embedding is computed."""

    SPACY = 1
    """[Spacy](https://spacy.io/models)."""

    TRANSFORMERS = 2
    """[HuggingFace Transformers](https://huggingface.co/models)."""

    SENTENCE_TRANSFORMERS = 3
    """
    [UKPLab Sentence
    Transformers](https://www.sbert.net/docs/pretrained_models.html)
    """

    TENSORFLOW_HUB = 4
    """
    Tensorflow Hub. Example: [Universal Sentence
    Encoder](https://tfhub.dev/google/universal-sentence-encoder/4).
    """

    OPENAI = 5
    """
    [OpenAI](https://platform.openai.com/docs/guides/embeddings/use-cases).
    """


@dataclass(eq=False, repr=False)
class NlpConfig(betterproto.Message):
    """Common message for configuring spacy."""

    language: str = betterproto.string_field(1)
    """
    Any language supported by spacy (e.g., `en`).
    [Reference](https://spacy.io/usage/models#languages).
    """

    spacy_model: str = betterproto.string_field(2)
    """
    Name of the trained spacy pipeline (e.g., `en_core_web_lg`). If empty, a
    blank spacy model will be used (e.g., if you only need embeddings and
    provide custom `embedding_models`. [Example: English
    models](https://spacy.io/models/en).
    """

    embedding_models: List["EmbeddingModel"] = betterproto.message_field(3)
    """
    List of embeddings to use for computing word/sentence vectors. If given,
    these embeddings will **override** the embeddings of the specified
    `spacy_model`. Multiple models are concatenated to each other, increasing
    the length of the resulting vector.
    """

    similarity_method: "SimilarityMethod" = betterproto.enum_field(4)
    """
    Mathematical function to determine a similarity score given two strings.
    """


@dataclass(eq=False, repr=False)
class SimilaritiesRequest(betterproto.Message):
    config: "NlpConfig" = betterproto.message_field(1)
    """Spacy config."""

    text_tuples: List["TextTuple"] = betterproto.message_field(2)
    """List of string pairs to compare."""

    extras: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(15)
    """Implementation-specific information can be encoded here"""


@dataclass(eq=False, repr=False)
class SimilaritiesResponse(betterproto.Message):
    similarities: List[float] = betterproto.double_field(1)
    """List of similarities ordered just like the original `text_tuples`."""

    extras: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(15)
    """Implementation-specific information can be encoded here"""


@dataclass(eq=False, repr=False)
class TextTuple(betterproto.Message):
    """Store a pair of strings."""

    text1: str = betterproto.string_field(1)
    text2: str = betterproto.string_field(2)


@dataclass(eq=False, repr=False)
class Strings(betterproto.Message):
    """Wrapper message to encode a list of strings that can also be `null`."""

    values: List[str] = betterproto.string_field(1)


@dataclass(eq=False, repr=False)
class DocBinRequest(betterproto.Message):
    config: "NlpConfig" = betterproto.message_field(1)
    """Spacy config."""

    texts: List[str] = betterproto.string_field(2)
    """List of strings to be processed."""

    attributes: Optional["Strings"] = betterproto.message_field(3, optional=True)
    """
    Attributes that shall be included in the DocBin object. Defaults to
    `("ORTH", "TAG", "HEAD", "DEP", "ENT_IOB", "ENT_TYPE", "ENT_KB_ID",
    "LEMMA", "MORPH", "POS")`. Possible values: `("IS_ALPHA", "IS_ASCII",
    "IS_DIGIT", "IS_LOWER", "IS_PUNCT", "IS_SPACE", "IS_TITLE", "IS_UPPER",
    "LIKE_URL", "LIKE_NUM", "LIKE_EMAIL", "IS_STOP", "IS_OOV_DEPRECATED",
    "IS_BRACKET", "IS_QUOTE", "IS_LEFT_PUNCT", "IS_RIGHT_PUNCT", "IS_CURRENCY",
    "ID", "ORTH", "LOWER", "NORM", "SHAPE", "PREFIX", "SUFFIX", "LENGTH",
    "CLUSTER", "LEMMA", "POS", "TAG", "DEP", "ENT_IOB", "ENT_TYPE", "ENT_ID",
    "ENT_KB_ID", "HEAD", "SENT_START", "SENT_END", "SPACY", "PROB", "LANG",
    "MORPH", "IDX")`. [Documentation](https://spacy.io/api/token#attributes).
    """

    enabled_pipes: "Strings" = betterproto.message_field(4, group="pipes")
    disabled_pipes: "Strings" = betterproto.message_field(5, group="pipes")
    embedding_levels: List["EmbeddingLevel"] = betterproto.enum_field(6)
    """
    List of vectors that shall be saved in the `DocBin` object. The computation
    is time-consuming, so you should only specify the embeddings you actually
    use!
    """

    extras: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(15)
    """Implementation-specific information can be encoded here"""


@dataclass(eq=False, repr=False)
class DocBinResponse(betterproto.Message):
    docbin: bytes = betterproto.bytes_field(1)
    """Serialized [`DocBin`](https://spacy.io/api/docbin) object"""

    extras: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(15)
    """Implementation-specific information can be encoded here"""


@dataclass(eq=False, repr=False)
class VectorsRequest(betterproto.Message):
    config: "NlpConfig" = betterproto.message_field(1)
    """Spacy config."""

    texts: List[str] = betterproto.string_field(2)
    """List of strings that shall be embedded (i.e., converted to vectors)."""

    embedding_levels: List["EmbeddingLevel"] = betterproto.enum_field(3)
    """
    List of vectors that shall be returned. The computation is time-consuming,
    so you should only specify the embeddings you actually use!
    """

    extras: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(15)
    """Implementation-specific information can be encoded here"""


@dataclass(eq=False, repr=False)
class VectorsResponse(betterproto.Message):
    vectors: List["VectorResponse"] = betterproto.message_field(1)
    """List of vectors whose order corresponds to the one of `texts`."""

    extras: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(15)
    """Implementation-specific information can be encoded here"""


@dataclass(eq=False, repr=False)
class VectorResponse(betterproto.Message):
    """
    Container object that includes vectors for all levels specified in
    `embedding_levels`.
    """

    document: "Vector" = betterproto.message_field(1)
    """One vector for the whole string."""

    tokens: List["Vector"] = betterproto.message_field(2)
    """Vectors for all tokens in the string."""

    sentences: List["Vector"] = betterproto.message_field(3)
    """Vectors for all sentences found in the string."""

    extras: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(15)
    """Implementation-specific information can be encoded here"""


@dataclass(eq=False, repr=False)
class Vector(betterproto.Message):
    """Container for storing a vector as a list of floats."""

    vector: List[float] = betterproto.double_field(1)


@dataclass(eq=False, repr=False)
class EmbeddingModel(betterproto.Message):
    """
    Specification of one model that is used to generate embeddings for strings.
    """

    model_type: "EmbeddingType" = betterproto.enum_field(1)
    """
    Each embedding has to be implemented, thus this enum is used to select the
    correct one.
    """

    model_name: str = betterproto.string_field(2)
    """
    You have to specify the name of the model that should be used by the
    selected impelemtation (i.e., `model_type`). We provide links to exemplary
    models for each implementation in the documentation of `EmbeddingType`.
    """

    pooling_type: "Pooling" = betterproto.enum_field(3, group="pooling")
    """Standard pooling functions like mean, min, max."""

    pmean: float = betterproto.double_field(4, group="pooling")
    """
    Power mean (or generalized mean). This method allows you to alter the
    computation of the mean representation. Special cases include arithmetic
    mean (p = 1), geometric mean (p = 0), harmonic mean (p = -1), minimum (p =
    -∞), maximum (p = ∞).
    [Wikipedia](https://en.wikipedia.org/wiki/Generalized_mean).
    [Paper](https://arxiv.org/abs/1803.01400).
    """


class NlpServiceStub(betterproto.ServiceStub):
    async def vectors(
        self,
        vectors_request: "VectorsRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "VectorsResponse":
        return await self._unary_unary(
            "/arg_services.nlp.v1.NlpService/Vectors",
            vectors_request,
            VectorsResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def similarities(
        self,
        similarities_request: "SimilaritiesRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "SimilaritiesResponse":
        return await self._unary_unary(
            "/arg_services.nlp.v1.NlpService/Similarities",
            similarities_request,
            SimilaritiesResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def doc_bin(
        self,
        doc_bin_request: "DocBinRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "DocBinResponse":
        return await self._unary_unary(
            "/arg_services.nlp.v1.NlpService/DocBin",
            doc_bin_request,
            DocBinResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )


class NlpServiceBase(ServiceBase):
    async def vectors(self, vectors_request: "VectorsRequest") -> "VectorsResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def similarities(
        self, similarities_request: "SimilaritiesRequest"
    ) -> "SimilaritiesResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def doc_bin(self, doc_bin_request: "DocBinRequest") -> "DocBinResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def __rpc_vectors(
        self, stream: "grpclib.server.Stream[VectorsRequest, VectorsResponse]"
    ) -> None:
        request = await stream.recv_message()
        response = await self.vectors(request)
        await stream.send_message(response)

    async def __rpc_similarities(
        self, stream: "grpclib.server.Stream[SimilaritiesRequest, SimilaritiesResponse]"
    ) -> None:
        request = await stream.recv_message()
        response = await self.similarities(request)
        await stream.send_message(response)

    async def __rpc_doc_bin(
        self, stream: "grpclib.server.Stream[DocBinRequest, DocBinResponse]"
    ) -> None:
        request = await stream.recv_message()
        response = await self.doc_bin(request)
        await stream.send_message(response)

    def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
        return {
            "/arg_services.nlp.v1.NlpService/Vectors": grpclib.const.Handler(
                self.__rpc_vectors,
                grpclib.const.Cardinality.UNARY_UNARY,
                VectorsRequest,
                VectorsResponse,
            ),
            "/arg_services.nlp.v1.NlpService/Similarities": grpclib.const.Handler(
                self.__rpc_similarities,
                grpclib.const.Cardinality.UNARY_UNARY,
                SimilaritiesRequest,
                SimilaritiesResponse,
            ),
            "/arg_services.nlp.v1.NlpService/DocBin": grpclib.const.Handler(
                self.__rpc_doc_bin,
                grpclib.const.Cardinality.UNARY_UNARY,
                DocBinRequest,
                DocBinResponse,
            ),
        }
