Wyvern makes it so much easier to define and auto retrieve features you need as well as connecting to your model, whether it’s a model that’s served from another service or a model that’s hosted locally within the wyvern pipeline service.

Wyvern has a ModelComponent which is the base component that takes in the inference request and a list of entities, and outputs a model output. The model output is a dictionary mapping entity identifiers to model outputs.

Let’s see how ModelComponent actually works

​
Model Inference

ModelComponent has an inference function which is the main entrance to model evaluation.

By default, the base ModelComponent slices entities into smaller batches and call batch_inference on each batch. The default batch size is 30. You should be able to configure the MODEL_BATCH_SIZE env variable to change the batch size.

In order to set up model inference, you only need to define a class that inherits ModelComponent and implement batch_inference.

You can also override this function if you want to customize the inference logic.

​
Register features for the model

Here’s an example of defining features you need for your model inference:

    @cached_property
    def manifest_feature_names(self) -> Set[str]:
        return {
            "RealtimeProductFeature:search_score",
            "RealtimeProductQueryFeature:query_category_similarity",
            "RealtimeProductQueryFeature:query_brand_similarity",
            "RealtimeProductQueryFeature:query_title_similarity",
            "RealtimeProductQueryFeature:query_description_similarity",
        }

The manifest_feature_names is a @cached_property (you can also just do @property but for cached_property gives a little performance advantage) and is a set of feature strings, with each string in the following format:

  • for Wyvern’s real-time features: realtime_feature_component_name + : + feature_name. The realtime_feature_component_name is the name of the RealtimeFeatureComponent you defined. By default, it is the name of your model class.
  • for batch features: FEATURE_VIEW + : + FEATURE_NAME

Wyvern’s framework will use features defined under manifest_feature_names

​
Define the whole model

​
Single Model

pipelines/product_ranking/models.py
import asyncio
from functools import cached_property
from typing import List, Set

from wyvern import (
    CompositeIdentifier,
    ModelComponent,
    ModelInput,
    ModelOutput,
    RankingRequest,
)

from pipelines.product_ranking.schemas import Product


class RankingModel(
    ModelComponent[
        ModelInput[
            Product,
            RankingRequest[Product],
        ],
        ModelOutput[float],
    ],
):
    @cached_property
    def manifest_feature_names(self) -> Set[str]:
        return {
            "RealtimeProductFeature:search_score",
            "RealtimeProductQueryFeature:query_category_similarity",
            "RealtimeProductQueryFeature:query_brand_similarity",
            "RealtimeProductQueryFeature:query_title_similarity",
            "RealtimeProductQueryFeature:query_description_similarity",
        }

    async def batch_inference(
        self,
        request: RankingRequest[Product],
        entities: List[Product],
        **kwargs,
    ) -> List[float]:
        return await asyncio.gather(
            *[self._inference_helper(request, entity) for entity in entities]
        )

    async def _inference_helper(
        self,
        request: RankingRequest[Product],
        entity: Product,
    ) -> float:
        score = 0.0
        composite_identifier = CompositeIdentifier(
            primary_identifier=entity.identifier,
            secondary_identifier=request.query.identifier,
        )
        search_score = self.get_feature(
            entity.identifier,
            "RealtimeProductFeature:search_score",
        )
        query_category_similarity = self.get_feature(
            composite_identifier,
            "RealtimeProductQueryFeature:query_category_similarity",
        )
        query_brand_similarity = self.get_feature(
            composite_identifier,
            "RealtimeProductQueryFeature:query_brand_similarity",
        )
        query_title_similarity = self.get_feature(
            composite_identifier,
            "RealtimeProductQueryFeature:query_title_similarity",
        )
        query_description_similarity = self.get_feature(
            composite_identifier,
            "RealtimeProductQueryFeature:query_description_similarity",
        )
        if search_score:
            score += 2 * search_score
        if query_category_similarity:
            score += 10 * query_category_similarity
        if query_brand_similarity:
            score += 20 * query_brand_similarity
        if query_title_similarity:
            score += 22 * query_title_similarity
        if query_description_similarity:
            score += 13 * query_description_similarity

        return score

As you could see in the example, the RankingModel inherits ModelComponent and defines manifest_feature_names and batch_inference. The batch_inference defines the behavior of model for batch inference. In this example of the _inference_helper, each inference is pretty much handling the formula we’ve talked about, which is:

model_score =
    2 * search_score
    + 10 * query_category_similarity
    + 20 * query_brand_similarity
    + 22 * query_title_similarity
    + 13 * query_description_similarity

​
ModelInput

This is the basic input format of a ModelComponent.

It contains the request information and the list of entities that’s related to the model

Reference: https://docs.wyvern.ai/sdk_ref#modelinput-objects

After the model is defined, now let’s go back to “Wyvern Pipeline” and see how to set up a pipeline in the code in the next page.

Was this page helpful?

ML Features In WyvernWyvern Pipeline