53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
|
import os
|
||
|
|
import logging
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
# from rpyc.utils.server import ThreadedServer
|
||
|
|
import typer
|
||
|
|
|
||
|
|
from ...utils.serve import ASRService
|
||
|
|
from plume.utils import lazy_callable
|
||
|
|
# from plume.models.wav2vec2_transformers.asr import Wav2Vec2TransformersASR
|
||
|
|
# from .asr import Wav2Vec2ASR
|
||
|
|
|
||
|
|
ThreadedServer = lazy_callable("rpyc.utils.server.ThreadedServer")
|
||
|
|
Wav2Vec2TransformersASR = lazy_callable(
|
||
|
|
"plume.models.wav2vec2_transformers.asr.Wav2Vec2TransformersASR"
|
||
|
|
)
|
||
|
|
|
||
|
|
app = typer.Typer()
|
||
|
|
|
||
|
|
|
||
|
|
@app.command()
|
||
|
|
def rpyc(
|
||
|
|
w2v_path: Path = "/path/to/base.pt",
|
||
|
|
ctc_path: Path = "/path/to/ctc.pt",
|
||
|
|
target_dict_path: Path = "/path/to/dict.ltr.txt",
|
||
|
|
port: int = int(os.environ.get("ASR_RPYC_PORT", "8044")),
|
||
|
|
):
|
||
|
|
w2vasr = Wav2Vec2TransformersASR(ctc_path, w2v_path, target_dict_path)
|
||
|
|
service = ASRService(w2vasr)
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.INFO,
|
||
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||
|
|
)
|
||
|
|
logging.info("starting asr server...")
|
||
|
|
t = ThreadedServer(service, port=port)
|
||
|
|
t.start()
|
||
|
|
|
||
|
|
|
||
|
|
@app.command()
|
||
|
|
def rpyc_dir(model_dir: Path, port: int = int(os.environ.get("ASR_RPYC_PORT", "8044"))):
|
||
|
|
ctc_path = model_dir / Path("ctc.pt")
|
||
|
|
w2v_path = model_dir / Path("base.pt")
|
||
|
|
target_dict_path = model_dir / Path("dict.ltr.txt")
|
||
|
|
rpyc(w2v_path, ctc_path, target_dict_path, port)
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
app()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|