Refactored code. Minor improvements

main.py

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