Skip to content

Commit

Permalink
Add auto for model/beam size
Browse files Browse the repository at this point in the history
  • Loading branch information
synesthesiam committed Dec 10, 2024
1 parent bfe6430 commit 85bd1e7
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 2 deletions.
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
# Changelog

## 2.4.0

- Add "auto" for model and beam size (0) to select values based on CPU

## 2.3.0

- Bump faster-whisper package to 1.1.0
Expand Down
2 changes: 1 addition & 1 deletion wyoming_faster_whisper/VERSION
Original file line number Diff line number Diff line change
@@ -1 +1 @@
2.3.0
2.4.0
16 changes: 15 additions & 1 deletion wyoming_faster_whisper/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import argparse
import asyncio
import logging
import platform
import re
from functools import partial

Expand All @@ -21,7 +22,7 @@ async def main() -> None:
parser.add_argument(
"--model",
required=True,
help="Name of faster-whisper model to use",
help="Name of faster-whisper model to use (or auto)",
)
parser.add_argument("--uri", required=True, help="unix:// or tcp://")
parser.add_argument(
Expand Down Expand Up @@ -52,6 +53,7 @@ async def main() -> None:
"--beam-size",
type=int,
default=5,
help="Size of beam during decoding (0 for auto)",
)
parser.add_argument(
"--initial-prompt",
Expand Down Expand Up @@ -79,6 +81,18 @@ async def main() -> None:
)
_LOGGER.debug(args)

# Automatic configuration for ARM
machine = platform.machine().lower()
is_arm = ("arm" in machine) or ("aarch" in machine)
if args.model == "auto":
args.model = "tiny-int8" if is_arm else "base-int8"
_LOGGER.debug("Model automatically selected: %s", args.model)

if args.beam_size <= 0:
args.beam_size = 1 if is_arm else 5
_LOGGER.debug("Beam size automatically selected: %s", args.beam_size)

# Resolve model name
model_name = args.model
match = re.match(r"^(tiny|base|small|medium)[.-]int8$", args.model)
if match:
Expand Down

0 comments on commit 85bd1e7

Please sign in to comment.