Parakeet-ROCm

Containerised, GPU-accelerated Automatic Speech Recognition (ASR) inference service for the NVIDIA Parakeet-TDT 0.6B v2 model, running on AMD ROCm GPUs.

Table of Contents

• Why This Project?
• Key Features
• What's Included
• Installation
  • Recommended: Docker Compose
  • Alternative: Local Development
• Configuration
• Usage
  • CLI
  • API Parameters
  • Output Files
• Development
• Contributing
• License

Why This Project?

This project bridges the gap between NVIDIA's cutting-edge ASR models and AMD GPU hardware through the ROCm platform. While NVIDIA's NeMo framework is primarily optimized for CUDA, this project enables running the powerful Parakeet-TDT model on AMD hardware with minimal performance impact.

Key Features

• High Accuracy: Leverages NVIDIA's state-of-the-art Parakeet-TDT 0.6B v2 model
• GPU Accelerated: Runs efficiently on AMD GPUs through ROCm platform
• Containerised: Dockerized deployment with ROCm support
• Multiple Formats: Export transcriptions in TXT, SRT, VTT, or JSON formats
• Timestamp Support: Word-level timestamps with intelligent segmentation
• Batch Processing: Process multiple files efficiently with configurable batch sizes
• Configurable: Environment-based configuration for all key parameters

What's Included

• CLI Tool: Typer-based command-line interface with rich progress tracking
• Modular Transcription Pipeline: parakeet_rocm/transcription split into CLI orchestration, per-file processing, and shared utilities
• Docker Support: Pre-configured container with ROCm, NeMo 2.2, and all dependencies
• Batch Processing: Efficient batch transcription with configurable chunking
• Multiple Output Formats: TXT, SRT, VTT, and JSON transcription outputs
• Timestamp Alignment: Word-level timestamp generation and intelligent subtitle segmentation
• GPU Acceleration: Optimized for AMD GPUs via ROCm platform

Installation

Recommended: Docker Compose

1. Clone the repository:

   git clone https://github.com/beecave-homelab/parakeet_rocm.git
   cd parakeet_rocm

2. Build the Docker image (first time ~10-15 min):

   pip install pdm
   pdm install -G rocm,webui
   # or: docker compose build

3. Run the container:

   parakeet-rocm --help
   # or: docker compose up

Alternative: Local Development

Prerequisites: Python 3.10, ROCm 6.4.1, PDM ≥2.15, ROCm PyTorch wheels in your --find-links.

1. Create lockfile and install dependencies (including ROCm extras):

   pdm install -G rocm,webui
   pip install requirements-all.txt # used as fallback for local development

2. Run unit tests:

   pytest -q

3. Transcribe a wav file locally:

   # Use the installed CLI script
   parakeet-rocm transcribe data/samples/voice_sample.wav

Configuration

The project uses environment variables for configuration. See .env.example for all options:

cp .env.example .env

Key configuration variables:

Variable	Description	Default
CHUNK_LEN_SEC	Audio chunk size in seconds for processing long files	300
BATCH_SIZE	Batch size for model inference	16
MAX_CPS	Max characters per second for subtitle readability	17
MIN_CPS	Min characters per second for subtitle readability	12
MAX_LINE_CHARS	Maximum characters per subtitle line	42

For ROCm-specific configuration, the following environment variables are set by default:

• PYTORCH_HIP_ALLOC_CONF=expandable_segments:True (mitigates ROCm memory fragmentation)
• HSA_OVERRIDE_GFX_VERSION=10.3.0 (required for some AMD GPUs)

Usage

CLI

The primary interface is a Typer-based CLI with rich help messages:

# Basic transcription
parakeet-rocm transcribe data/samples/sample.wav

# Transcribe multiple files
parakeet-rocm transcribe file1.wav file2.wav

# Specify output directory and format
parakeet-rocm transcribe --output-dir ./transcripts --output-format srt file.wav

# Adjust batch size for performance
parakeet-rocm transcribe --batch-size 8 file.wav

# Enable word-level timestamps
parakeet-rocm transcribe --word-timestamps file.wav

# Refine timestamps using stable-ts with VAD
parakeet-rocm transcribe --word-timestamps --stabilize --vad file.wav

# Continuous directory watching (auto-transcribe new files)
parakeet-rocm transcribe --watch data/watch/ --verbose

# Get help
parakeet-rocm --help
parakeet-rocm transcribe --help

API Parameters

Parameter	Type	Description	Default
audio_files	List[pathlib.Path]	Path to one or more audio files to transcribe	required
--model	str	Hugging Face Hub or local path to the NeMo ASR model	nvidia/parakeet-tdt-0.6b-v2
--output-dir	pathlib.Path	Directory to save the transcription outputs	./output
--output-format	str	Format for the output file(s) (txt, srt, vtt, json)	txt
--batch-size	int	Batch size for transcription inference	16 (from env)
--chunk-len-sec	int	Segment length in seconds for chunked transcription	300 (from env)
--word-timestamps	bool	Enable word-level timestamp generation	False
--stabilize	bool	Refine word timestamps using stable-ts	False
--vad	bool	Enable VAD during stabilization	False
--demucs	bool	Use Demucs denoiser during stabilization	False
--vad-threshold	float	VAD probability threshold	0.35
--watch	List[str]	Watch directory or wildcard pattern(s) for new audio/video files	None
--overwrite	bool	Overwrite existing output files	False
--verbose	bool	Enable verbose output (shows detailed logs from NeMo and Transformers)	False
--quiet	bool	Suppress console output except progress bar (Note: Logs are now suppressed by default unless --verbose is used)	False
--no-progress	bool	Disable the Rich progress bar while still showing created file paths	False
--fp16	bool	Enable half-precision (FP16) inference	False

Output Files

Transcriptions are saved in the specified output directory with filenames based on the input files. Supported formats include:

• TXT: Plain text transcription
• SRT: SubRip subtitle format with timestamps
• VTT: Web Video Text Tracks format
• JSON: Structured output with detailed timing information

Development

See project-overview.md for complete architecture and developer documentation.

For local development:

1. Install development dependencies:

   pdm install -G rocm -G dev

2. Run tests:

   pytest -q

3. Code formatting and linting:

   bash scripts/clean_codebase.sh

Contributing

Pull requests are welcome. For major changes, please open an issue first to discuss your proposal.

See project-overview.md for complete architecture and developer documentation.

License

This project is licensed under the MIT License - see the LICENSE file for details.