diff --git a/main.py b/main.py index 56279f5..6a305d5 100644 --- a/main.py +++ b/main.py @@ -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" - -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", - torch_dtype=torch.float16, - device_map="auto" -).to(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() - - 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() - - total_samples = waveform.shape[0] - chunk_size = chunk_length_s * SAMPLING_FREQUENCY - num_chunks = (total_samples + chunk_size - 1) // chunk_size - - logger.info(f"File length - {total_samples/SAMPLING_FREQUENCY:.1f} seconds, splitting on {num_chunks} chunks by {chunk_length_s} seconds") - - transcripts = [] - - for i in range(num_chunks): - start = i * chunk_size - end = min((i + 1) * chunk_size, total_samples) - chunk = waveform[start:end].cpu().numpy() - - inputs = processor(chunk, sampling_rate=SAMPLING_FREQUENCY, return_tensors="pt") - input_features = inputs.input_features.to(device).to(torch.float16) - - with torch.no_grad(): - predicted_ids = model.generate( - input_features, - language=language, - task="transcribe", - temperature=0.0 - ) - - text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0] - transcripts.append(text) - - logger.info(f"Чанк {i+1}/{num_chunks} готов ({(end/SAMPLING_FREQUENCY):.1f} сек)") - - end_time = time.time() - logger.info(f"Transcription completed - {end_time - start_time:.2f} seconds") - - return " ".join(transcripts) - - except Exception as e: - logger.error(f"Transcription error: {str(e)}") - raise - - -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 +class AudioTranscription: + model_name = "openai/whisper-large-v2" - except Exception as e: - logger.error(f"Error while splitting to chunks: {str(e)}") - raise + 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(self.device) + + logger.info("Model loaded.") + + self.waveform, self.sampling_rate = torchaudio.load(filepath, format="mp3", backend="ffmpeg") + logger.info(f"Successfully loaded file {filepath}.") + + except Exception as e: + logger.error(f"Unable to load file {self.filepath}: {e}") + raise + + def resample(self) -> None: + self.waveform = torchaudio.functional.resample(self.waveform, self.sampling_rate, 16000) + + 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) + + def split_to_chunks(self, chunk_length_s: int = 30) -> None: + logger.info(f"Splitting audio on chunks...") + + 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 + + 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) -chunks = split_into_chunks(f"{FILENAME}.mp3") + with torch.no_grad(): + predicted_ids = self.model.generate( + input_features, + language=self.language, + task="transcribe", + temperature=0.0 + ) + + text = self.processor.batch_decode(predicted_ids, skip_special_tokens=True)[0] + + 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 in {end_time - start_time:.2f} seconds") + + except Exception as e: + logger.error(f"Errors occured while processing chunks: {e}") + + 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}") \ No newline at end of file