ソースを参照

add language parameter

subDesTagesMitExtraKaese 2 年 前
コミット
f2183cfc3c
3 ファイル変更10 行追加3 行削除
  1. 1 0
      README.md
  2. 1 1
      main.py
  3. 8 2
      speech_recognition.py

+ 1 - 0
README.md

@@ -23,6 +23,7 @@ services:
       - "USERNAME=@stt-bot:example.com"
       - "PASSWORD=<password>"
       - "ASR_MODEL=tiny"
+      - "ASR_LANGUAGE=en"
 ```
 
 ## Configuration

+ 1 - 1
main.py

@@ -23,7 +23,7 @@ config.ignore_unverified_devices = True
 config.store_path = '/data/crypto_store/'
 bot = botlib.Bot(creds, config)
 
-asr = ASR(os.getenv('ASR_MODEL', 'tiny'))
+asr = ASR(os.getenv('ASR_MODEL', 'tiny'), os.getenv('ASR_LANGUAGE', 'en'))
 
 @bot.listener.on_custom_event(nio.RoomMessageAudio)
 async def on_message_audio(room, event):

+ 8 - 2
speech_recognition.py

@@ -22,10 +22,12 @@ def convert_audio(data: bytes) -> bytes:
 MODELS = ["tiny.en", "tiny", "base.en", "base", "small.en", "small", "medium.en", "medium", "large"]
 
 class ASR():
-  def __init__(self, model = "tiny"):
+  def __init__(self, model = "tiny", language = "en"):
     if model not in MODELS:
       raise ValueError(f"Invalid model: {model}. Must be one of {MODELS}")
     self.model = model
+    self.language = language
+
     if not os.path.exists("/data/models"):
       os.mkdir("/data/models")
     self.model_path = f"/data/models/ggml-{model}.bin"
@@ -42,7 +44,11 @@ class ASR():
     async with self.lock:
       convert_audio(audio)
       proc = await asyncio.create_subprocess_exec(
-          "./main", "-m", self.model_path, "-f", "audio.wav", "--no_timestamps", 
+          "./main",
+          "-m", self.model_path,
+          "-l", self.language,
+          "-f", "audio.wav",
+          "--no_timestamps", 
           stdout=asyncio.subprocess.PIPE,
           stderr=asyncio.subprocess.PIPE
         )