Eloquence-Backend / utils /transcription.py
PMS61
server files added
ebf88d6
raw
history blame
2.86 kB
import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from pydub import AudioSegment
import soundfile as sf
import os
model_name = "openai/whisper-base"
processor = WhisperProcessor.from_pretrained(model_name)
model = WhisperForConditionalGeneration.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
def preprocess_audio(input_audio_path, output_audio_path):
"""
Converts audio to 16kHz WAV format.
Args:
input_audio_path (str): Path to the input audio file.
output_audio_path (str): Path to save the processed audio file.
Returns:
str: Path to the processed audio file.
"""
audio = AudioSegment.from_file(input_audio_path)
audio = audio.set_frame_rate(16000).set_channels(1)
audio.export(output_audio_path, format="wav")
return output_audio_path
def split_audio(audio_path, chunk_length_ms=30000):
"""
Splits audio into chunks of specified length.
Args:
audio_path (str): Path to the audio file.
chunk_length_ms (int): Length of each chunk in milliseconds.
Returns:
list: List of audio chunks.
"""
audio = AudioSegment.from_file(audio_path)
chunks = [audio[i : i + chunk_length_ms] for i in range(0, len(audio), chunk_length_ms)]
return chunks
def transcribe_chunk(audio_chunk, chunk_index):
"""
Transcribes a single audio chunk.
Args:
audio_chunk (AudioSegment): The audio chunk to transcribe.
chunk_index (int): Index of the chunk.
Returns:
str: Transcription of the chunk.
"""
temp_path = f"temp_chunk_{chunk_index}.wav"
audio_chunk.export(temp_path, format="wav")
audio, sampling_rate = sf.read(temp_path)
inputs = processor(audio, sampling_rate=16000, return_tensors="pt")
input_features = inputs.input_features.to(device)
predicted_ids = model.generate(input_features)
transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
os.remove(temp_path) # Clean up temporary file
return transcription
def speech_to_text_long(audio_path):
"""
Transcribes a long audio file by splitting it into chunks.
Args:
audio_path (str): Path to the audio file.
Returns:
str: Full transcription of the audio.
"""
processed_audio_path = "processed_audio.wav"
preprocess_audio(audio_path, processed_audio_path)
# Split audio into chunks
chunks = split_audio(processed_audio_path, chunk_length_ms=30000) # 30 seconds per chunk
transcriptions = []
for idx, chunk in enumerate(chunks):
print(f"Transcribing chunk {idx + 1} of {len(chunks)}...")
transcription = transcribe_chunk(chunk, idx)
transcriptions.append(transcription)
return " ".join(transcriptions)