mirror of
https://github.com/shirayu/whispering.git
synced 2024-05-08 19:33:09 +00:00
88 lines
2.1 KiB
Python
88 lines
2.1 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import sys
|
|
from typing import Final, List, Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
from pydantic import BaseModel, Field, root_validator
|
|
from whisper.audio import N_FRAMES
|
|
|
|
|
|
class WhisperConfig(BaseModel):
|
|
model_name: str
|
|
device: str
|
|
language: str
|
|
fp16: bool = True
|
|
|
|
@root_validator
|
|
def validate_model_name(cls, values):
|
|
if values["model_name"].endswith(".en") and values["language"] not in {
|
|
"en",
|
|
"English",
|
|
}:
|
|
raise ValueError("English only model")
|
|
return values
|
|
|
|
|
|
CURRENT_PROTOCOL_VERSION: Final[int] = int("000_006_003")
|
|
|
|
|
|
class Context(BaseModel, arbitrary_types_allowed=True):
|
|
protocol_version: int
|
|
timestamp: float = 0.0
|
|
buffer_tokens: List[torch.Tensor] = []
|
|
buffer_mel: Optional[torch.Tensor] = None
|
|
nosoeech_skip_count: Optional[int] = None
|
|
|
|
temperatures: List[float]
|
|
patience: Optional[float] = None
|
|
compression_ratio_threshold: Optional[float] = 2.4
|
|
logprob_threshold: Optional[float] = -1.0
|
|
no_captions_threshold: Optional[float] = 0.6
|
|
best_of: int = 5
|
|
beam_size: Optional[int] = None
|
|
no_speech_threshold: Optional[float] = 0.6
|
|
logprob_threshold: Optional[float] = -1.0
|
|
compression_ratio_threshold: Optional[float] = 2.4
|
|
buffer_threshold: Optional[float] = 0.5
|
|
vad_threshold: float
|
|
max_nospeech_skip: int
|
|
mel_frame_min_num: int = Field(N_FRAMES, ge=1, le=N_FRAMES)
|
|
|
|
data_type: str = "float32"
|
|
|
|
|
|
class ParsedChunk(BaseModel):
|
|
start: float
|
|
end: float
|
|
text: str
|
|
tokens: List[int]
|
|
temperature: float
|
|
avg_logprob: float
|
|
compression_ratio: float
|
|
no_speech_prob: float
|
|
|
|
|
|
class SpeechSegment(BaseModel, arbitrary_types_allowed=True):
|
|
start_block_idx: int
|
|
end_block_idx: int
|
|
audio: np.ndarray
|
|
|
|
|
|
class StdoutWriter:
|
|
def open(self, *args, **kwargs):
|
|
return self
|
|
|
|
def __enter__(self, *args, **kwargs):
|
|
return self
|
|
|
|
def __exit__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def flush(self, *args, **kwargs):
|
|
sys.stdout.flush()
|
|
|
|
def write(self, text, *args, **kwargs):
|
|
sys.stdout.write(text)
|