336 lines
9.8 KiB
Python
336 lines
9.8 KiB
Python
import tqdm
|
|
import json
|
|
from fastapi.responses import StreamingResponse
|
|
import wave
|
|
import pydub
|
|
import io
|
|
import hashlib
|
|
import argparse
|
|
import uvicorn
|
|
from typing import Annotated, Any, BinaryIO, Literal, Generator, Tuple, Iterable
|
|
from fastapi import (
|
|
File,
|
|
HTTPException,
|
|
Query,
|
|
UploadFile,
|
|
Form,
|
|
FastAPI,
|
|
Request,
|
|
WebSocket,
|
|
)
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from src.whisper_ctranslate2.whisper_ctranslate2 import Transcribe
|
|
from src.whisper_ctranslate2.writers import format_timestamp
|
|
from faster_whisper.transcribe import Segment, TranscriptionInfo
|
|
import opencc
|
|
from prometheus_fastapi_instrumentator import Instrumentator
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="0.0.0.0", type=str)
|
|
parser.add_argument("--port", default=5000, type=int)
|
|
parser.add_argument("--model", default="large-v3", type=str)
|
|
parser.add_argument("--device", default="auto", type=str)
|
|
parser.add_argument("--cache_dir", default=None, type=str)
|
|
parser.add_argument("--local_files_only", default=False, type=bool)
|
|
parser.add_argument("--threads", default=4, type=int)
|
|
args = parser.parse_args()
|
|
app = FastAPI()
|
|
# Instrument your app with default metrics and expose the metrics
|
|
Instrumentator().instrument(app).expose(app, endpoint="/konele/metrics")
|
|
ccc = opencc.OpenCC("t2s.json")
|
|
|
|
print("Loading model...")
|
|
transcriber = Transcribe(
|
|
model_path=args.model,
|
|
device=args.device,
|
|
device_index=0,
|
|
compute_type="default",
|
|
threads=args.threads,
|
|
cache_directory=args.cache_dir,
|
|
local_files_only=args.local_files_only,
|
|
)
|
|
print("Model loaded!")
|
|
|
|
|
|
# allow all cors
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
def stream_writer(generator: Generator[dict[str, Any], Any, None]):
|
|
for segment in generator:
|
|
yield "data: " + json.dumps(segment, ensure_ascii=False) + "\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
|
|
def text_writer(generator: Generator[dict[str, Any], Any, None]):
|
|
for segment in generator:
|
|
yield segment["text"].strip() + "\n"
|
|
|
|
|
|
def tsv_writer(generator: Generator[dict[str, Any], Any, None]):
|
|
yield "start\tend\ttext\n"
|
|
for i, segment in enumerate(generator):
|
|
start_time = str(round(1000 * segment["start"]))
|
|
end_time = str(round(1000 * segment["end"]))
|
|
text = segment["text"]
|
|
yield f"{start_time}\t{end_time}\t{text}\n"
|
|
|
|
|
|
def srt_writer(generator: Generator[dict[str, Any], Any, None]):
|
|
for i, segment in enumerate(generator):
|
|
start_time = format_timestamp(
|
|
segment["start"], decimal_marker=",", always_include_hours=True
|
|
)
|
|
end_time = format_timestamp(
|
|
segment["end"], decimal_marker=",", always_include_hours=True
|
|
)
|
|
text = segment["text"]
|
|
yield f"{i}\n{start_time} --> {end_time}\n{text}\n\n"
|
|
|
|
|
|
def vtt_writer(generator: Generator[dict[str, Any], Any, None]):
|
|
yield "WEBVTT\n\n"
|
|
for i, segment in enumerate(generator):
|
|
start_time = format_timestamp(segment["start"])
|
|
end_time = format_timestamp(segment["end"])
|
|
text = segment["text"]
|
|
yield f"{start_time} --> {end_time}\n{text}\n\n"
|
|
|
|
|
|
def build_json_result(
|
|
generator: Iterable[Segment],
|
|
info: dict,
|
|
) -> dict[str, Any]:
|
|
segments = [i for i in generator]
|
|
return info | {
|
|
"text": "\n".join(i["text"] for i in segments),
|
|
"segments": segments,
|
|
}
|
|
|
|
|
|
def stream_builder(
|
|
audio: BinaryIO,
|
|
task: str,
|
|
vad_filter: bool,
|
|
language: str | None,
|
|
initial_prompt: str = "",
|
|
repetition_penalty: float = 1.0,
|
|
) -> Tuple[Iterable[dict], dict]:
|
|
segments, info = transcriber.model.transcribe(
|
|
audio=audio,
|
|
language=language,
|
|
task=task,
|
|
initial_prompt=initial_prompt,
|
|
word_timestamps=True,
|
|
repetition_penalty=repetition_penalty,
|
|
)
|
|
print(
|
|
"Detected language '%s' with probability %f"
|
|
% (info.language, info.language_probability)
|
|
)
|
|
def wrap():
|
|
last_pos = 0
|
|
with tqdm.tqdm(total=info.duration, unit="seconds", disable=True) as pbar:
|
|
for segment in segments:
|
|
start, end, text = segment.start, segment.end, segment.text
|
|
pbar.update(end - last_pos)
|
|
last_pos = end
|
|
data = segment._asdict()
|
|
if data.get('words') is not None:
|
|
data["words"] = [i._asdict() for i in data["words"]]
|
|
if info.language == "zh":
|
|
data["text"] = ccc.convert(data["text"])
|
|
yield data
|
|
|
|
info_dict = info._asdict()
|
|
if info_dict['transcription_options'] is not None:
|
|
info_dict['transcription_options'] = info_dict['transcription_options']._asdict()
|
|
if info_dict['vad_options'] is not None:
|
|
info_dict['vad_options'] = info_dict['vad_options']._asdict()
|
|
|
|
return wrap(), info_dict
|
|
|
|
|
|
@app.websocket("/k6nele/status")
|
|
@app.websocket("/konele/status")
|
|
async def konele_status(
|
|
websocket: WebSocket,
|
|
):
|
|
await websocket.accept()
|
|
await websocket.send_json(dict(num_workers_available=1))
|
|
await websocket.close()
|
|
|
|
|
|
@app.websocket("/k6nele/ws")
|
|
@app.websocket("/konele/ws")
|
|
async def konele_ws(
|
|
websocket: WebSocket,
|
|
task: Literal["transcribe", "translate"] = "transcribe",
|
|
lang: str = "und",
|
|
initial_prompt: str = "",
|
|
vad_filter: bool = False,
|
|
content_type: Annotated[str, Query(alias="content-type")] = "audio/x-raw",
|
|
):
|
|
await websocket.accept()
|
|
|
|
# convert lang code format (eg. en-US to en)
|
|
lang = lang.split("-")[0]
|
|
|
|
data = b""
|
|
while True:
|
|
try:
|
|
data += await websocket.receive_bytes()
|
|
if data[-3:] == b"EOS":
|
|
break
|
|
except:
|
|
break
|
|
|
|
md5 = hashlib.md5(data).hexdigest()
|
|
|
|
# create fake file for wave.open
|
|
file_obj = io.BytesIO()
|
|
|
|
if content_type.startswith("audio/x-flac"):
|
|
pydub.AudioSegment.from_file(io.BytesIO(data), format="flac").export(
|
|
file_obj, format="wav"
|
|
)
|
|
else:
|
|
buffer = wave.open(file_obj, "wb")
|
|
buffer.setnchannels(1)
|
|
buffer.setsampwidth(2)
|
|
buffer.setframerate(16000)
|
|
buffer.writeframes(data)
|
|
|
|
file_obj.seek(0)
|
|
|
|
generator, info = stream_builder(
|
|
audio=file_obj,
|
|
task=task,
|
|
vad_filter=vad_filter,
|
|
language=None if lang == "und" else lang,
|
|
initial_prompt=initial_prompt,
|
|
)
|
|
result = build_json_result(generator, info)
|
|
|
|
text = result.get("text", "")
|
|
|
|
await websocket.send_json(
|
|
{
|
|
"status": 0,
|
|
"segment": 0,
|
|
"result": {"hypotheses": [{"transcript": text}], "final": True},
|
|
"id": md5,
|
|
}
|
|
)
|
|
await websocket.close()
|
|
|
|
|
|
@app.post("/k6nele/post")
|
|
@app.post("/konele/post")
|
|
async def translateapi(
|
|
request: Request,
|
|
task: Literal["transcribe", "translate"] = "transcribe",
|
|
lang: str = "und",
|
|
initial_prompt: str = "",
|
|
vad_filter: bool = False,
|
|
):
|
|
content_type = request.headers.get("Content-Type", "")
|
|
|
|
# convert lang code format (eg. en-US to en)
|
|
lang = lang.split("-")[0]
|
|
|
|
splited = [i.strip() for i in content_type.split(",") if "=" in i]
|
|
info = {k: v for k, v in (i.split("=") for i in splited)}
|
|
|
|
channels = int(info.get("channels", "1"))
|
|
rate = int(info.get("rate", "16000"))
|
|
|
|
body = await request.body()
|
|
md5 = hashlib.md5(body).hexdigest()
|
|
|
|
# create fake file for wave.open
|
|
file_obj = io.BytesIO()
|
|
|
|
if content_type.startswith("audio/x-flac"):
|
|
pydub.AudioSegment.from_file(io.BytesIO(body), format="flac").export(
|
|
file_obj, format="wav"
|
|
)
|
|
else:
|
|
buffer = wave.open(file_obj, "wb")
|
|
buffer.setnchannels(channels)
|
|
buffer.setsampwidth(2)
|
|
buffer.setframerate(rate)
|
|
buffer.writeframes(body)
|
|
|
|
file_obj.seek(0)
|
|
|
|
generator, info = stream_builder(
|
|
audio=file_obj,
|
|
task=task,
|
|
vad_filter=vad_filter,
|
|
language=None if lang == "und" else lang,
|
|
initial_prompt=initial_prompt,
|
|
)
|
|
result = build_json_result(generator, info)
|
|
|
|
text = result.get("text", "")
|
|
|
|
return {
|
|
"status": 0,
|
|
"hypotheses": [{"utterance": text}],
|
|
"id": md5,
|
|
}
|
|
|
|
|
|
@app.post("/v1/audio/transcriptions")
|
|
async def transcription(
|
|
file: UploadFile = File(...),
|
|
prompt: str = Form(""),
|
|
response_format: str = Form("json"),
|
|
task: str = Form("transcribe"),
|
|
language: str = Form("und"),
|
|
vad_filter: bool = Form(False),
|
|
repetition_penalty: float = Form(1.0),
|
|
):
|
|
"""Transcription endpoint
|
|
|
|
User upload audio file in multipart/form-data format and receive transcription in response
|
|
"""
|
|
|
|
# timestamp as filename, keep original extension
|
|
generator, info = stream_builder(
|
|
audio=io.BytesIO(file.file.read()),
|
|
task=task,
|
|
vad_filter=vad_filter,
|
|
language=None if language == "und" else language,
|
|
repetition_penalty=repetition_penalty,
|
|
)
|
|
|
|
# special function for streaming response (OpenAI API does not have this)
|
|
if response_format == "stream":
|
|
return StreamingResponse(
|
|
stream_writer(generator),
|
|
media_type="text/event-stream",
|
|
)
|
|
elif response_format == "json":
|
|
return build_json_result(generator, info)
|
|
elif response_format == "text":
|
|
return StreamingResponse(text_writer(generator), media_type="text/plain")
|
|
elif response_format == "tsv":
|
|
return StreamingResponse(tsv_writer(generator), media_type="text/plain")
|
|
elif response_format == "srt":
|
|
return StreamingResponse(srt_writer(generator), media_type="text/plain")
|
|
elif response_format == "vtt":
|
|
return StreamingResponse(vtt_writer(generator), media_type="text/plain")
|
|
|
|
raise HTTPException(400, "Invailed response_format")
|
|
|
|
|
|
uvicorn.run(app, host=args.host, port=args.port)
|