[AI] stt 를 구현해보자. (feat. whisper)

늘 공부하는 괴짜·2025년 5월 31일

AI : STT & TTS

목록 보기

1/1

1. 일단 pyenv 먼저 설치

이 글을 쓰기전에 리허설 과정에서 버전 문제가 계속 나와 진행이 불가하였다.
오랜만에 pyenv 를 설치해본다.

% brew install pyenv
% pyenv install 3.10.13

1-1. stt 디렉토리에 버전 지정

% cd stt
% pyenv local 3.10.13

1-2. uv 환경 설정

% uv venv .venv
% source .venv/bin/activate
% python --version  # Python 3.10.13

2. ffmpeg 설치

% brew install ffmpeg

3. 필요 패키지 설치

% uv pip install torch torchaudio pyannote.audio transformers

4. convert.py

화자분리 모델은 pyannote/speaker-diarization
stt 모델은 openai/whisper-large-v2 를 사용하였다.
https://huggingface.co/openai/whisper-large-v2

import os
import subprocess
import torchaudio
import torch
from pyannote.audio import Pipeline
from transformers import WhisperProcessor, WhisperForConditionalGeneration

# Hugging Face 토큰 (환경 변수로 설정되어 있어야 함)
hf_token = os.getenv("HUGGINGFACE_API_KEY")

input_audio = "./stt_1.wav"
preprocessed_audio = "./stt_preprocessed.wav"

# 1. ffmpeg로 16kHz mono, 노이즈 제거된 WAV 생성
subprocess.run([
    "ffmpeg", "-y",
    "-i", input_audio,
    "-ac", "1",
    "-ar", "16000",
    "-af", "afftdn=nf=-35",
    preprocessed_audio
], check=True)

# 2. pyannote 화자 분리
# num_speakers를 지정하지 않으면 모델이 자동으로 화자 수를 추정합니다.
# 만약 화자 수를 정확히 알고 있다면, num_speakers=N 과 같이 지정할 수 있습니다.
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hf_token)
diarization = pipeline(preprocessed_audio)

# 3. Whisper 모델과 프로세서 로드
processor = WhisperProcessor.from_pretrained("openai/whisper-large-v2", use_auth_token=hf_token)
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large-v2", use_auth_token=hf_token)

# ✅ MPS 장치 설정 (M1 / M2 / M3)
if torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model.to(device)
model.eval()


# 4. 샘플링 레이트 조회
info = torchaudio.info(preprocessed_audio)
sr = info.sample_rate

# 최소 대화 지속 시간 설정 (초 단위)
# 이 값보다 짧은 대화는 무시됩니다.
min_duration_threshold = 0.2 # 예시: 0.5초 미만 대화는 생략

# 5. diarization 결과 구간별 STT 수행
for turn, _, speaker in diarization.itertracks(yield_label=True):
    duration = turn.end - turn.start

    # 설정된 최소 지속 시간보다 짧은 대화는 건너뛰기
    if duration < min_duration_threshold:
        continue

    start_frame = int(turn.start * sr)
    num_frames = int(duration * sr)

    waveform, sr_loaded = torchaudio.load(preprocessed_audio, frame_offset=start_frame, num_frames=num_frames)
    assert sr == sr_loaded, "샘플링 레이트 불일치!"

    inputs = processor(waveform.squeeze(), sampling_rate=16000, return_tensors="pt")
    input_features = inputs.input_features.to(device)

    generated_ids = model.generate(input_features, language="ko", task="transcribe")
    transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

    print(f"[{speaker}] {turn.start:.1f}s ~ {turn.end:.1f}s: {transcription}")