Olive_Whisper_ASR / infer_server.py
sam2ai's picture
Synced repo using 'sync_with_huggingface' Github Action
28ebae2
raw
history blame
4.99 kB
import argparse
import asyncio
import functools
import json
import os
from io import BytesIO
import uvicorn
from fastapi import FastAPI, BackgroundTasks, File, Body, UploadFile, Request
from fastapi.responses import StreamingResponse
from faster_whisper import WhisperModel
from starlette.staticfiles import StaticFiles
from starlette.templating import Jinja2Templates
from zhconv import convert
from utils.data_utils import remove_punctuation
from utils.utils import add_arguments, print_arguments
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
add_arg("host", type=str, default="0.0.0.0", help="")
add_arg("port", type=int, default=5000, help="")
add_arg("model_path", type=str, default="models/sam2ai/whisper-odia-small-finetune-int8-ct2", help="")
add_arg("use_gpu", type=bool, default=False, help="")
add_arg("use_int8", type=bool, default=True, help="")
add_arg("beam_size", type=int, default=10, help="")
add_arg("num_workers", type=int, default=2, help="")
add_arg("vad_filter", type=bool, default=True, help="")
add_arg("local_files_only", type=bool, default=True, help="")
args = parser.parse_args()
print_arguments(args)
#
assert os.path.exists(args.model_path), f"{args.model_path}"
#
if args.use_gpu:
if not args.use_int8:
model = WhisperModel(args.model_path, device="cuda", compute_type="float16",
num_workers=args.num_workers, local_files_only=args.local_files_only)
else:
model = WhisperModel(args.model_path, device="cuda",
compute_type="int8_float16", num_workers=args.num_workers,
local_files_only=args.local_files_only)
else:
model = WhisperModel(args.model_path, device="cpu",
compute_type="int8", num_workers=args.num_workers,
local_files_only=args.local_files_only)
#
# _, _ = model.transcribe("dataset/test.wav", beam_size=5)
app = FastAPI(title="OdiaGenAI Whisper ASR")
app.mount('/static', StaticFiles(directory='static'), name='static')
templates = Jinja2Templates(directory="templates")
model_semaphore = None
def release_model_semaphore():
model_semaphore.release()
def recognition(file: File, to_simple: int,
remove_pun: int, language: str = "bn",
task: str = "transcribe"
):
segments, info = model.transcribe(file, beam_size=10, task=task, language=language, vad_filter=args.vad_filter)
for segment in segments:
text = segment.text
if to_simple == 1:
# text = convert(text, '')
pass
if remove_pun == 1:
# text = remove_punctuation(text)
pass
ret = {"result": text, "start": round(segment.start, 2), "end": round(segment.end, 2)}
#
yield json.dumps(ret).encode() + b"\0"
@app.post("/recognition_stream")
async def api_recognition_stream(
to_simple: int = Body(1, description="", embed=True),
remove_pun: int = Body(0, description="", embed=True),
language: str = Body("bn", description="", embed=True),
task: str = Body("transcribe", description="", embed=True),
audio: UploadFile = File(..., description="")
):
global model_semaphore
if language == "None": language = None
if model_semaphore is None:
model_semaphore = asyncio.Semaphore(5)
await model_semaphore.acquire()
contents = await audio.read()
data = BytesIO(contents)
generator = recognition(
file=data, to_simple=to_simple,
remove_pun=remove_pun, language=language,
task=task
)
background_tasks = BackgroundTasks()
background_tasks.add_task(release_model_semaphore)
return StreamingResponse(generator, background=background_tasks)
@app.post("/recognition")
async def api_recognition(
to_simple: int = Body(1, description="", embed=True),
remove_pun: int = Body(0, description="", embed=True),
language: str = Body("bn", description="", embed=True),
task: str = Body("transcribe", description="", embed=True),
audio: UploadFile = File(..., description="")
):
if language == "None":language=None
contents = await audio.read()
data = BytesIO(contents)
generator = recognition(
file=data, to_simple=to_simple,
remove_pun=remove_pun, language=language,
task=task
)
results = []
for output in generator:
output = json.loads(output[:-1].decode("utf-8"))
results.append(output)
ret = {"results": results, "code": 0}
return ret
@app.get("/")
async def index(request: Request):
return templates.TemplateResponse(
"index.html", {"request": request, "id": id}
)
if __name__ == '__main__':
uvicorn.run(app, host=args.host, port=args.port)