inference_types.py

from dataclasses import dataclass
from typing import List, Optional


@dataclass
class Message:
    role: str
    content: str


@dataclass
class GenerateRequest:
    messages: List[Message]
    max_output_length: Optional[int] = 2048
    temperature: Optional[float] = 0.1
    top_p: Optional[float] = 1
    top_k: Optional[int] = 16
    repetition_penalty: Optional[float] = 1.1
    streaming: Optional[bool] = False


@dataclass
class GenerateResponse:
    response: Message
    generation_tokens: int
    generation_time: float
    tokens_per_second: float


@dataclass
class BatchGenerateRequest:
    batch_messages: List[List[Message]]
    max_output_length: Optional[int] = 2048
    temperature: Optional[float] = 0.1
    top_p: Optional[float] = 1
    top_k: Optional[int] = 16
    repetition_penalty: Optional[float] = 1.1
    streaming: Optional[bool] = False


@dataclass
class BatchGenerateResponse:
    batch_responses: List[List[Message]]
    batch_generation_tokens: List[int]
    generation_time: float
    max_throughput_tokens_per_second: float