# generated by datamodel-codegen:
#   filename:  topk_by_token.yaml

from __future__ import annotations

from typing import List, Optional

from pydantic import BaseModel, Field


class NPActivationTopkByTokenRequest(BaseModel):
    secret: str = Field(..., description='API secret for the inference server')
    prompt: str = Field(..., description='Input text prompt to get activations for')
    model: str = Field(..., description='Name of the model to test activations on')
    source: str = Field(
        ...,
        description='Source identifier - could be an SAE ID (eg 5-gemmascope-res-16k). Must be specified with "index", or not at NPActivationAllRequest.',
    )
    top_k: Optional[int] = Field(None, description='The number of features to include for each token position.')
    ignore_bos: bool = Field(
        ..., description='Whether or not to include features whose highest activation value is the BOS token.'
    )


class TopFeature(BaseModel):
    feature_index: int = Field(..., description='The index of the feature in the SAE.')
    activation_value: float = Field(..., description='The activation value of this feature at this token position.')


class NPActivationTopkByTokenResult(BaseModel):
    token_position: int = Field(..., description='The index of the token in the prompt.')
    token: str = Field(..., description='The token string')
    top_features: List[TopFeature]


class NPActivationTopkByTokenResponse(BaseModel):
    results: List[NPActivationTopkByTokenResult]
    tokens: List[str]
