Refactored code. Minor improvements
This commit is contained in:
@@ -3,10 +3,7 @@ import torch
|
||||
import torchaudio
|
||||
import logging
|
||||
import time
|
||||
import asyncio
|
||||
|
||||
FILENAME = "sample"
|
||||
SAMPLING_FREQUENCY = 16000
|
||||
from tqdm import tqdm
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
@@ -23,104 +20,122 @@ if torch.cuda.is_available():
|
||||
print(f"Name of GPU: {torch.cuda.get_device_name(0)}")
|
||||
print("=== Check completed ===")
|
||||
|
||||
device = "cuda"
|
||||
class AudioTranscription:
|
||||
model_name = "openai/whisper-large-v2"
|
||||
|
||||
logger.info("Loading model WhisperProcessor...")
|
||||
processor = WhisperProcessor.from_pretrained("openai/whisper-large-v2")
|
||||
logger.info("Loading model WhisperForConditionalGeneration...")
|
||||
model = WhisperForConditionalGeneration.from_pretrained(
|
||||
"openai/whisper-large-v2",
|
||||
filepath: str
|
||||
waveform: torch.Tensor
|
||||
sampling_rate: int
|
||||
|
||||
chunks: list = []
|
||||
# one chunk size in seconds
|
||||
chunk_size: int
|
||||
device = "cuda"
|
||||
processor: WhisperProcessor
|
||||
model: WhisperForConditionalGeneration
|
||||
# self.device here and all previous shit
|
||||
|
||||
language = "ru"
|
||||
|
||||
all_transcription: str = ""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
filepath: str,
|
||||
language = "ru",
|
||||
device = "cuda",
|
||||
model_name = "openai/whisper-large-v2"
|
||||
) -> None:
|
||||
self.filepath = filepath
|
||||
self.language = language
|
||||
self.device = device
|
||||
self.model_name = model_name
|
||||
self.chunks: list = []
|
||||
try:
|
||||
logger.info("Loading model WhisperProcessor...")
|
||||
self.processor = WhisperProcessor.from_pretrained(self.model_name)
|
||||
|
||||
self.model = WhisperForConditionalGeneration.from_pretrained(
|
||||
self.model_name,
|
||||
torch_dtype=torch.float16,
|
||||
device_map="auto"
|
||||
).to(device)
|
||||
logger.info("Model loaded")
|
||||
).to(self.device)
|
||||
|
||||
logger.info("Model loaded.")
|
||||
|
||||
def transcribe_long(path: str, language="ru", chunk_length_s: int = 30):
|
||||
logger.info(f"Starting transcription of long file: {path}")
|
||||
start_time = time.time()
|
||||
self.waveform, self.sampling_rate = torchaudio.load(filepath, format="mp3", backend="ffmpeg")
|
||||
logger.info(f"Successfully loaded file {filepath}.")
|
||||
|
||||
try:
|
||||
waveform, sr = torchaudio.load(path, format="mp3", backend="ffmpeg")
|
||||
if waveform.shape[0] > 1:
|
||||
waveform = waveform.mean(dim=0, keepdim=True)
|
||||
waveform = torchaudio.functional.resample(waveform, sr, SAMPLING_FREQUENCY).squeeze()
|
||||
except Exception as e:
|
||||
logger.error(f"Unable to load file {self.filepath}: {e}")
|
||||
raise
|
||||
|
||||
total_samples = waveform.shape[0]
|
||||
chunk_size = chunk_length_s * SAMPLING_FREQUENCY
|
||||
num_chunks = (total_samples + chunk_size - 1) // chunk_size
|
||||
def resample(self) -> None:
|
||||
self.waveform = torchaudio.functional.resample(self.waveform, self.sampling_rate, 16000)
|
||||
|
||||
logger.info(f"File length - {total_samples/SAMPLING_FREQUENCY:.1f} seconds, splitting on {num_chunks} chunks by {chunk_length_s} seconds")
|
||||
def to_mono(self):
|
||||
if self.waveform.shape[0] > 1:
|
||||
self.waveform = self.waveform.mean(dim=0, keepdim=True)
|
||||
self.waveform = self.waveform.squeeze(0)
|
||||
|
||||
transcripts = []
|
||||
def split_to_chunks(self, chunk_length_s: int = 30) -> None:
|
||||
logger.info(f"Splitting audio on chunks...")
|
||||
|
||||
for i in range(num_chunks):
|
||||
start = i * chunk_size
|
||||
end = min((i + 1) * chunk_size, total_samples)
|
||||
chunk = waveform[start:end].cpu().numpy()
|
||||
self.chunk_size = chunk_length_s * 16000 # 16kHz after resampling
|
||||
total_samples = self.waveform.shape[0]
|
||||
chunks_count = (total_samples + self.chunk_size - 1) // self.chunk_size
|
||||
|
||||
inputs = processor(chunk, sampling_rate=SAMPLING_FREQUENCY, return_tensors="pt")
|
||||
input_features = inputs.input_features.to(device).to(torch.float16)
|
||||
logger.info(f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {chunk_length_s} seconds per chunk.")
|
||||
|
||||
self.chunks = []
|
||||
for idx in range(chunks_count):
|
||||
start = idx * self.chunk_size
|
||||
end = min((idx + 1) * self.chunk_size, total_samples)
|
||||
chunk = self.waveform[start:end].cpu().numpy().astype("float32")
|
||||
self.chunks.append(chunk)
|
||||
|
||||
def process_chunk(
|
||||
self,
|
||||
chunk
|
||||
) -> str:
|
||||
inputs = self.processor(chunk, sampling_rate=16000, return_tensors="pt")
|
||||
input_features = inputs.input_features.to(self.device).to(torch.float16)
|
||||
|
||||
with torch.no_grad():
|
||||
predicted_ids = model.generate(
|
||||
predicted_ids = self.model.generate(
|
||||
input_features,
|
||||
language=language,
|
||||
language=self.language,
|
||||
task="transcribe",
|
||||
temperature=0.0
|
||||
)
|
||||
|
||||
text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
|
||||
transcripts.append(text)
|
||||
text = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
|
||||
|
||||
logger.info(f"Чанк {i+1}/{num_chunks} готов ({(end/SAMPLING_FREQUENCY):.1f} сек)")
|
||||
return text
|
||||
|
||||
def process_all_chunks(self) -> None:
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
for chunk_idx in tqdm(range(len(self.chunks))):
|
||||
# remake without strings (slow asf)
|
||||
self.all_transcription += " " + self.process_chunk(self.chunks[chunk_idx])
|
||||
|
||||
end_time = time.time()
|
||||
logger.info(f"Transcription completed - {end_time - start_time:.2f} seconds")
|
||||
|
||||
return " ".join(transcripts)
|
||||
logger.info(f"Transcription completed in {end_time - start_time:.2f} seconds")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Transcription error: {str(e)}")
|
||||
raise
|
||||
logger.error(f"Errors occured while processing chunks: {e}")
|
||||
|
||||
|
||||
def split_into_chunks(
|
||||
filepath: str,
|
||||
chunk_length_s: int = 30
|
||||
) -> list:
|
||||
try:
|
||||
chunks = []
|
||||
waveform, sr = torchaudio.load(filepath, format="mp3", backend="ffmpeg")
|
||||
|
||||
if waveform.shape[0] > 1:
|
||||
waveform = waveform.mean(dim=0, keepdim=True)
|
||||
|
||||
logger.info(f"Started splitting file into chunks with length {chunk_length_s}")
|
||||
total_samples = waveform.shape[0]
|
||||
chunk_size = chunk_length_s
|
||||
|
||||
num_chunks = (total_samples + chunk_size - 1) // chunk_size
|
||||
logger.info(f"File length - {total_samples/SAMPLING_FREQUENCY:.1f} seconds, разобьём на {num_chunks} чанков по {chunk_length_s} секунд")
|
||||
|
||||
for i in range(num_chunks):
|
||||
start = i * chunk_size
|
||||
end = min((i + 1) * chunk_size, total_samples)
|
||||
chunk = waveform[start:end].cpu().numpy()
|
||||
|
||||
chunks.append(chunk)
|
||||
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error while splitting to chunks: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
chunks = split_into_chunks(f"{FILENAME}.mp3")
|
||||
def transcribe_audio(self) -> str:
|
||||
self.resample()
|
||||
self.to_mono()
|
||||
self.split_to_chunks()
|
||||
self.process_all_chunks()
|
||||
return self.all_transcription
|
||||
|
||||
try:
|
||||
result = transcribe_long(f"{FILENAME}.mp3")
|
||||
print(result)
|
||||
track = AudioTranscription("sample.mp3")
|
||||
print(track.transcribe_audio())
|
||||
except Exception as e:
|
||||
logger.error(f"Execution error: {e}")
|
||||
Reference in New Issue
Block a user