mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 23:51:00 +02:00
Vladimir Mandic
66 lines
1.6 KiB
Python
66 lines
1.6 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, List, Optional
|
|
from transformers import StoppingCriteria
|
|
|
|
|
|
@dataclass
|
|
class Model:
|
|
name: Optional[str] = None
|
|
cls: str = "UnknownModelClass"
|
|
tokenizer: str = "UnknownTokenizerClass"
|
|
processor: Optional[str] = None
|
|
type: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
max_context_tokens: int = 4096
|
|
max_new_tokens: int = 512
|
|
stream: bool = False
|
|
temperature: float = 0.2
|
|
top_p: float = 0.9
|
|
top_k: int = 50
|
|
repetition_penalty: float = 1.0
|
|
use_cache: bool = True
|
|
|
|
|
|
@dataclass
|
|
class Attention:
|
|
implementation: str = "eager"
|
|
flash: bool = False
|
|
arch: str = ""
|
|
|
|
|
|
@dataclass
|
|
class Req:
|
|
id: int
|
|
client: Optional[str] = None
|
|
url: Optional[str] = None
|
|
config: dict[str, Any] = None
|
|
|
|
|
|
class CustomTokenStopCriteria(StoppingCriteria):
|
|
def __init__(self, stop_ids: List[int]):
|
|
super().__init__()
|
|
self.stop_ids = stop_ids
|
|
|
|
def __call__(self, input_ids: Any, scores: Any, **kwargs) -> bool:
|
|
return input_ids[0][-1].item() in self.stop_ids
|
|
|
|
|
|
class Stats:
|
|
id: int = 0
|
|
messages: int = 0
|
|
prompt: int = 0
|
|
ttft: float = 0.0
|
|
latency: float = 0.0
|
|
tps: float = 0.0
|
|
tools: int = 0
|
|
tokens: dict[str, int] = {}
|
|
chunks: dict[str, int] = {}
|
|
streaming: bool = None
|
|
thinking: bool = None
|
|
|
|
def __str__(self):
|
|
return f"Res(id={self.id}, messages={self.messages}, prompt={self.prompt}, ttft={self.ttft:.3f}, latency={self.latency:.3f}, tps={self.tps:.3f}, tools={self.tools}, tokens={self.tokens}, chunks={self.chunks}, streaming={self.streaming}, thinking={self.thinking})"
|