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

from __future__ import annotations

from typing import List, Optional

from pydantic import BaseModel, Field


class NPGetSaeVectorRequest(BaseModel):
    model: str = Field(..., example='gpt2-small')
    source: str = Field(..., example='0-res-jb')
    index: int = Field(..., example=14057)
    secret: str = Field(..., description='Authentication secret for the API', example='your-secret-key')


class NPGetSaeVectorResponse(BaseModel):
    vector: List[float]


class NPFeature(BaseModel):
    model: str = Field(..., example='gpt2-small')
    source: str = Field(..., example='0-res-jb')
    index: int = Field(..., example=14057)


class TopkDecoderCossimFeature(BaseModel):
    feature: Optional[NPFeature] = None
    cosine_similarity: Optional[float] = None


class NPSaeTopkByDecoderCosSimResponse(BaseModel):
    feature: Optional[NPFeature] = None
    topk_decoder_cossim_features: Optional[List[TopkDecoderCossimFeature]] = None


class NPSaeTopkByDecoderCosSimRequest(BaseModel):
    feature: Optional[NPFeature] = None
    secret: str
    num_results: int
