add: option of task and lang

This commit is contained in:
2023-11-15 16:19:20 +08:00
parent bb52b0c44e
commit 62b7431655

View File

@@ -3,7 +3,7 @@ import io
import hashlib import hashlib
import argparse import argparse
import uvicorn import uvicorn
from typing import Any from typing import Any, Literal
from fastapi import File, UploadFile, Form, FastAPI, Request, WebSocket, Response from fastapi import File, UploadFile, Form, FastAPI, Request, WebSocket, Response
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from src.whisper_ctranslate2.whisper_ctranslate2 import Transcribe, TranscriptionOptions from src.whisper_ctranslate2.whisper_ctranslate2 import Transcribe, TranscriptionOptions
@@ -112,6 +112,7 @@ def get_options(*, initial_prompt=""):
@app.websocket("/konele/ws") @app.websocket("/konele/ws")
async def konele_ws( async def konele_ws(
websocket: WebSocket, websocket: WebSocket,
task: Literal["transcribe", "translate"] = "transcribe",
lang: str = "und", lang: str = "und",
): ):
await websocket.accept() await websocket.accept()
@@ -143,9 +144,8 @@ async def konele_ws(
result = transcriber.inference( result = transcriber.inference(
audio=file_obj, audio=file_obj,
# Enter translate mode if target language is English task=task,
task="translate" if lang == "en-US" else "transcribe", language=lang if lang != "und" else None, # type: ignore
language=None, # type: ignore
verbose=False, verbose=False,
live=False, live=False,
options=options, options=options,
@@ -168,6 +168,7 @@ async def konele_ws(
@app.post("/konele/post") @app.post("/konele/post")
async def translateapi( async def translateapi(
request: Request, request: Request,
task: Literal["transcribe", "translate"] = "transcribe",
lang: str = "und", lang: str = "und",
): ):
content_type = request.headers.get("Content-Type", "") content_type = request.headers.get("Content-Type", "")
@@ -196,9 +197,8 @@ async def translateapi(
result = transcriber.inference( result = transcriber.inference(
audio=file_obj, audio=file_obj,
# Enter translate mode if target language is English task=task,
task="translate" if lang == "en-US" else "transcribe", language=lang if lang != "und" else None, # type: ignore
language=None, # type: ignore
verbose=False, verbose=False,
live=False, live=False,
options=options, options=options,