feat: refactor Wav2Vec2 character ASR to support quantization and improve model loading

src/AI_Models/wave2vec_inference.py

@@ -0,0 +1,183 @@
+import torch
+from transformers import AutoModelForCTC, AutoProcessor, Wav2Vec2Processor
+import onnxruntime as rt
+import numpy as np
+import librosa
+
+
+class Wave2Vec2Inference:
+    def __init__(self, model_name, hotwords=[], use_lm_if_possible=True, use_gpu=True):
+        self.device = "cpu"
+        if use_lm_if_possible:
+            self.processor = AutoProcessor.from_pretrained(model_name)
+        else:
+            self.processor = Wav2Vec2Processor.from_pretrained(model_name)
+        self.model = AutoModelForCTC.from_pretrained(model_name)
+        self.model.to(self.device)
+        self.hotwords = hotwords
+        self.use_lm_if_possible = use_lm_if_possible
+
+    def buffer_to_text(self, audio_buffer):
+        if len(audio_buffer) == 0:
+            return ""
+
+        inputs = self.processor(
+            torch.tensor(audio_buffer),
+            sampling_rate=16_000,
+            return_tensors="pt",
+            padding=True,
+        )
+
+        with torch.no_grad():
+            logits = self.model(
+                inputs.input_values.to(self.device),
+                attention_mask=inputs.attention_mask.to(self.device),
+            ).logits
+
+        if hasattr(self.processor, "decoder") and self.use_lm_if_possible:
+            transcription = self.processor.decode(
+                logits[0].cpu().numpy(),
+                hotwords=self.hotwords,
+                # hotword_weight=self.hotword_weight,
+                output_word_offsets=True,
+            )
+            confidence = transcription.lm_score / len(transcription.text.split(" "))
+            transcription: str = transcription.text
+        else:
+            predicted_ids = torch.argmax(logits, dim=-1)
+            transcription: str = self.processor.batch_decode(predicted_ids)[0]
+            # confidence = self.confidence_score(logits, predicted_ids)
+        return transcription.lower()
+
+    def confidence_score(self, logits, predicted_ids):
+        scores = torch.nn.functional.softmax(logits, dim=-1)
+        pred_scores = scores.gather(-1, predicted_ids.unsqueeze(-1))[:, :, 0]
+        mask = torch.logical_and(
+            predicted_ids.not_equal(self.processor.tokenizer.word_delimiter_token_id),
+            predicted_ids.not_equal(self.processor.tokenizer.pad_token_id),
+        )
+
+        character_scores = pred_scores.masked_select(mask)
+        total_average = torch.sum(character_scores) / len(character_scores)
+        return total_average
+
+    def file_to_text(self, filename):
+        import librosa
+
+        audio_input, samplerate = librosa.load(filename, sr=16000)
+        return self.buffer_to_text(audio_input)
+
+
+class Wave2Vec2ONNXInference:
+    def __init__(self, model_name, onnx_path):
+        self.processor = Wav2Vec2Processor.from_pretrained(model_name)
+        # self.model = Wav2Vec2ForCTC.from_pretrained(model_name)
+        options = rt.SessionOptions()
+        options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
+        self.model = rt.InferenceSession(onnx_path, options)
+
+    def buffer_to_text(self, audio_buffer):
+        if len(audio_buffer) == 0:
+            return ""
+
+        inputs = self.processor(
+            torch.tensor(audio_buffer),
+            sampling_rate=16_000,
+            return_tensors="np",
+            padding=True,
+        )
+
+        input_values = inputs.input_values
+        onnx_outputs = self.model.run(
+            None, {self.model.get_inputs()[0].name: input_values}
+        )[0]
+        prediction = np.argmax(onnx_outputs, axis=-1)
+
+        transcription = self.processor.decode(prediction.squeeze().tolist())
+        return transcription.lower()
+
+    def file_to_text(self, filename):
+        audio_input, samplerate = librosa.load(filename, sr=16000)
+        return self.buffer_to_text(audio_input)
+
+
+from onnxruntime.quantization.quantize import quantize
+from transformers import Wav2Vec2ForCTC
+import torch
+import argparse
+
+# took that script from: https://github.com/ccoreilly/wav2vec2-service/blob/master/convert_torch_to_onnx.py
+
+
+def convert_to_onnx(model_id_or_path, onnx_model_name):
+    print(f"Converting {model_id_or_path} to onnx")
+    model = Wav2Vec2ForCTC.from_pretrained(model_id_or_path)
+    audio_len = 250000
+
+    x = torch.randn(1, audio_len, requires_grad=True)
+
+    torch.onnx.export(
+        model,  # model being run
+        x,  # model input (or a tuple for multiple inputs)
+        onnx_model_name,  # where to save the model (can be a file or file-like object)
+        export_params=True,  # store the trained parameter weights inside the model file
+        opset_version=14,  # the ONNX version to export the model to
+        do_constant_folding=True,  # whether to execute constant folding for optimization
+        input_names=["input"],  # the model's input names
+        output_names=["output"],  # the model's output names
+        dynamic_axes={
+            "input": {1: "audio_len"},  # variable length axes
+            "output": {1: "audio_len"},
+        },
+    )
+
+
+def quantize_onnx_model(onnx_model_path, quantized_model_path):
+    print("Starting quantization...")
+    from onnxruntime.quantization import quantize_dynamic, QuantType
+
+    quantize_dynamic(
+        onnx_model_path, quantized_model_path, weight_type=QuantType.QUInt8
+    )
+
+    print(f"Quantized model saved to: {quantized_model_path}")
+
+
+def export_to_onnx(
+    model: str = "facebook/wav2vec2-large-960h-lv60-self", quantize: bool = False
+):
+    onnx_model_name = model.split("/")[-1] + ".onnx"
+    convert_to_onnx(model, onnx_model_name)
+    if quantize:
+        quantized_model_name = model.split("/")[-1] + ".quant.onnx"
+        quantize_onnx_model(onnx_model_name, quantized_model_name)
+
+
+if __name__ == "__main__":
+    from loguru import logger
+    import time
+
+    asr = Wave2Vec2Inference("facebook/wav2vec2-large-960h-lv60-self")
+
+    # Warm up runs
+    print("Warming up...")
+    for i in range(2):
+        asr.file_to_text("test.wav")
+        print(f"Warm up {i+1} completed")
+
+    # Test runs
+    print("Running tests...")
+    times = []
+    for i in range(10):
+        start_time = time.time()
+        text = asr.file_to_text("test.wav")
+        end_time = time.time()
+        execution_time = end_time - start_time
+        times.append(execution_time)
+        print(f"Test {i+1}: {execution_time:.3f}s - {text}")
+
+    # Calculate average time
+    average_time = sum(times) / len(times)
+    print(f"\nAverage execution time: {average_time:.3f}s")
+    print(f"Min time: {min(times):.3f}s")
+    print(f"Max time: {max(times):.3f}s")

src/apis/controllers/speaking_controller.py

@@ -1,9 +1,4 @@
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException, APIRouter
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from typing import List, Dict, Optional
-import tempfile
-import os
+from typing import List, Dict
 import numpy as np
 import librosa
 import nltk
@@ -11,14 +6,15 @@ import eng_to_ipa as ipa
 import torch
 import re
 from collections import defaultdict
-import warnings
-from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
 from loguru import logger
-import onnxruntime
 import time
-
+from src.AI_Models.wave2vec_inference import (
+    Wave2Vec2Inference,
+    Wave2Vec2ONNXInference,
+    export_to_onnx,
+)
 
 # Download required NLTK data
 try:
@@ -128,7 +124,7 @@ class Wav2Vec2CharacterASR:
         self,
         model_name: str = "facebook/wav2vec2-large-960h-lv60-self",
         onnx: bool = False,
-        onnx_model_path: str = "./wav2vec2_asr.onnx",
+        quantized: bool = False,
     ):
         """
         Initialize Wav2Vec2 character-level model
@@ -138,185 +134,48 @@ class Wav2Vec2CharacterASR:
             onnx: If True, use ONNX runtime for inference. If False, use Transformers
             onnx_model_path: Path to the ONNX model file (only used if onnx=True)
         """
-        self.model_name = model_name
         self.use_onnx = onnx
-        self.onnx_model_path = onnx_model_path
         self.sample_rate = 16000
-
-        print(f"Loading Wav2Vec2 character model: {model_name}")
-        print(f"Using {'ONNX' if onnx else 'Transformers'} for inference")
-
-        if self.use_onnx:
-            self._init_onnx_model()
-        else:
-            self._init_transformers_model()
-
-    def _init_onnx_model(self):
-        """Initialize ONNX model and processor"""
-        # Check if ONNX model exists, if not create it
-        if not os.path.exists(self.onnx_model_path):
-            print(f"ONNX model not found at {self.onnx_model_path}. Creating it...")
-            self._create_onnx_model()
-
-        try:
-            # Load ONNX model
-            self.session = onnxruntime.InferenceSession(self.onnx_model_path)
-            self.input_name = self.session.get_inputs()[0].name
-            self.output_name = self.session.get_outputs()[0].name
-
-            # Load processor
-            self.processor = Wav2Vec2Processor.from_pretrained(self.model_name)
-
-            print("ONNX Wav2Vec2 character model loaded successfully")
-
-        except Exception as e:
-            print(f"Error loading ONNX model: {e}")
-            raise
-
-    def _init_transformers_model(self):
-        """Initialize Transformers model and processor"""
-        try:
-            self.processor = Wav2Vec2Processor.from_pretrained(self.model_name)
-            self.model = Wav2Vec2ForCTC.from_pretrained(self.model_name)
-            self.model.eval()
-            print("Wav2Vec2 character model loaded successfully")
-        except Exception as e:
-            print(f"Error loading model {self.model_name}: {e}")
-            # Fallback to base model
-            fallback_model = "facebook/wav2vec2-base-960h"
-            print(f"Trying fallback model: {fallback_model}")
-            try:
-                self.processor = Wav2Vec2Processor.from_pretrained(fallback_model)
-                self.model = Wav2Vec2ForCTC.from_pretrained(fallback_model)
-                self.model.eval()
-                self.model_name = fallback_model
-                print("Fallback model loaded successfully")
-            except Exception as e2:
-                raise Exception(
-                    f"Failed to load both models. Original error: {e}, Fallback error: {e2}"
-                )
-
-    def _create_onnx_model(self):
-        """Create ONNX model if it doesn't exist"""
-        try:
-            # Import the converter from model_convert
-            from src.model_convert.wav2vec2onnx import Wav2Vec2ONNXConverter
-
-            print("Creating new ONNX model...")
-            converter = Wav2Vec2ONNXConverter(self.model_name)
-            created_path = converter.convert_to_onnx(
-                onnx_path=self.onnx_model_path,
-                input_length=160000,  # 10 seconds
-                opset_version=14,
+        self.model_name = model_name
+        # Check thử path của onnx model có tồn tại hay không
+        if onnx:
+            import os
+
+            if not os.path.exists(
+                "wav2vec2-large-960h-lv60-self"
+                + (".quant" if quantized else "")
+                + ".onnx"
+            ):
+
+                export_to_onnx(model_name, quantize=quantized)
+        self.model = (
+            Wave2Vec2Inference(model_name)
+            if not onnx
+            else Wave2Vec2ONNXInference(
+                model_name,
+                "wav2vec2-large-960h-lv60-self"
+                + (".quant" if quantized else "")
+                + ".onnx",
             )
-            print(f"✓ ONNX model created successfully at: {created_path}")
-
-        except ImportError as e:
-            print(f"Error importing Wav2Vec2ONNXConverter: {e}")
-            raise e
+        )
 
     def transcribe_to_characters(self, audio_path: str) -> Dict:
-        """
-        Transcribe audio directly to characters (no language model correction)
-        Returns raw character sequence as produced by the model
-        """
-        if self.use_onnx:
-            return self._transcribe_onnx(audio_path)
-        else:
-            return self._transcribe_transformers(audio_path)
-
-    def _transcribe_onnx(self, audio_path: str) -> Dict:
-        """Transcribe using ONNX runtime"""
-        try:
-            # Load audio
-            start_time = time.time()
-            speech, sr = librosa.load(audio_path, sr=self.sample_rate)
-
-            # Prepare input for ONNX
-            input_values = self.processor(
-                speech, sampling_rate=self.sample_rate, return_tensors="np"
-            ).input_values
-
-            # Run ONNX inference
-            ort_inputs = {self.input_name: input_values}
-            ort_outputs = self.session.run([self.output_name], ort_inputs)
-            logits = ort_outputs[0]
-
-            # Get predictions
-            predicted_ids = np.argmax(logits, axis=-1)
-
-            # Decode to characters directly
-            character_transcript = self.processor.batch_decode(predicted_ids)[0]
-            logger.info(f"character_transcript {character_transcript}")
-
-            # Clean up character transcript
-            character_transcript = self._clean_character_transcript(
-                character_transcript
-            )
-
-            # Convert characters to phoneme-like representation
-            phoneme_like_transcript = self._characters_to_phoneme_representation(
-                character_transcript
-            )
-
-            # Calculate confidence scores
-            confidence_scores = self._calculate_confidence_scores(logits)
-            logger.info(
-                f"Wav2Vec2 ONNX transcription time: {time.time() - start_time:.2f}s"
-            )
-
-            return {
-                "character_transcript": character_transcript,
-                "phoneme_representation": phoneme_like_transcript,
-                "raw_predicted_ids": predicted_ids[0].tolist(),
-                "confidence_scores": confidence_scores[:100],  # Limit for JSON
-            }
-
-        except Exception as e:
-            print(f"ONNX transcription error: {e}")
-            return self._empty_result()
-
-    def _transcribe_transformers(self, audio_path: str) -> Dict:
-        """Transcribe using Transformers"""
         try:
-            # Load audio
             start_time = time.time()
-            speech, sr = librosa.load(audio_path, sr=self.sample_rate)
-
-            # Prepare input
-            input_values = self.processor(
-                speech, sampling_rate=self.sample_rate, return_tensors="pt"
-            ).input_values
-
-            # Get model predictions (no language model involved)
-            with torch.no_grad():
-                logits = self.model(input_values).logits
-                predicted_ids = torch.argmax(logits, dim=-1)
-
-            # Decode to characters directly
-            character_transcript = self.processor.batch_decode(predicted_ids)[0]
-
-            # Clean up character transcript
+            character_transcript = self.model.file_to_text(audio_path)
             character_transcript = self._clean_character_transcript(
                 character_transcript
             )
 
-            # Convert characters to phoneme-like representation
             phoneme_like_transcript = self._characters_to_phoneme_representation(
                 character_transcript
             )
 
-            logger.info(
-                f"Transformers transcription time: {time.time() - start_time:.2f}s"
-            )
+            logger.info(f"Transcription time: {time.time() - start_time:.2f}s")
 
             return {
                 "character_transcript": character_transcript,
                 "phoneme_representation": phoneme_like_transcript,
-                "raw_predicted_ids": predicted_ids[0].tolist(),
-                "confidence_scores": torch.softmax(logits, dim=-1)
-                .max(dim=-1)[0][0]
-                .tolist()[:100],  # Limit for JSON
             }
 
         except Exception as e:
@@ -988,7 +847,7 @@ class SimplePronunciationAssessor:
         if mode == "advanced":
             print("Step 1: Using Wav2Vec2 character transcription...")
             asr_result = self.wav2vec2_asr.transcribe_to_characters(audio_path)
-            model_info = f"Wav2Vec2-Character ({self.wav2vec2_asr.model_name})"
+            model_info = f"Wav2Vec2-Character ({self.wav2vec2_asr.model})"
         else:  # normal mode
             print("Step 1: Using Whisper transcription...")
             asr_result = self.whisper_asr.transcribe_to_text(audio_path)
@@ -1046,19 +905,3 @@ class SimplePronunciationAssessor:
 
         total_score = sum(comparison["score"] for comparison in phoneme_comparisons)
         return total_score / len(phoneme_comparisons)
-
-
-def convert_numpy_types(obj):
-    """Convert numpy types to Python native types"""
-    if isinstance(obj, np.integer):
-        return int(obj)
-    elif isinstance(obj, np.floating):
-        return float(obj)
-    elif isinstance(obj, np.ndarray):
-        return obj.tolist()
-    elif isinstance(obj, dict):
-        return {key: convert_numpy_types(value) for key, value in obj.items()}
-    elif isinstance(obj, list):
-        return [convert_numpy_types(item) for item in obj]
-    else:
-        return obj

src/apis/routes/speaking_route.py

@@ -10,8 +10,8 @@ from src.apis.controllers.speaking_controller import (
     SimpleG2P,
     PhonemeComparator,
     SimplePronunciationAssessor,
-    convert_numpy_types,
 )
+from src.utils.speaking_utils import convert_numpy_types
 
 warnings.filterwarnings("ignore")