optimize finetuned model

app/services/processor.py

@@ -12,16 +12,22 @@ import librosa
 import torch
 
 from app.core.config import get_settings
-from app.services.transcription import TranscriptionService,  WordTimestamp
+from app.services.transcription import TranscriptionService
 from app.services.diarization import DiarizationService, SpeakerSegment, DiarizationResult
-from app.services.alignment import AlignmentService
-from app.schemas.models import TranscriptSegment
-
 
 logger = logging.getLogger(__name__)
 settings = get_settings()
 
 
+@dataclass
+class TranscriptSegment:
+    """A transcribed segment with speaker info."""
+    start: float
+    end: float
+    speaker: str
+    role: Optional[str]
+    text: str
+
 
 @dataclass
 class ProcessingResult:
@@ -36,45 +42,76 @@ class ProcessingResult:
     txt_content: str = ""
     csv_content: str = ""
 
-
-
-def pad_and_refine_tensor(
-    waveform: torch.Tensor,
-    sr: int,
-    start_s: float,
-    end_s: float,
-    pad_ms: int = 200,
-) -> Tuple[float, float]:
-
-    total_len = waveform.shape[1]
-    s = max(int((start_s - pad_ms / 1000) * sr), 0)
-    e = min(int((end_s + pad_ms / 1000) * sr), total_len)
-
-    if e <= s:
-        return start_s, end_s
-
-    return s / sr, e / sr
-
-
-def infer_roles_from_diarization(
+def assign_speaker_to_word(
+    word_start: float,
     diarization_segments: List[SpeakerSegment],
-) -> Dict[str, str]:
-
-    dur = defaultdict(float)
-
-    for s in diarization_segments:
-        dur[s.speaker] += s.end - s.start
-
-    if not dur:
-        return {}
-
-    agent = max(dur, key=dur.get)
+) -> str:
+    """
+    Assign speaker to word using diarization.
+    """
+    for seg in diarization_segments:
+        if seg.start <= word_start <= seg.end:
+            return seg.speaker
+
+    # fallback: nearest diar segment
+    return min(
+        diarization_segments,
+        key=lambda s: abs((s.start + s.end) / 2 - word_start)
+    ).speaker
+
+
+def group_words_into_segments(
+    words: List[dict],
+    diarization_segments: List[SpeakerSegment],
+    speaker_map: Dict[str, str],
+    roles: Dict[str, str],
+    max_word_gap_s: float = 0.6,
+) -> List[TranscriptSegment]:
+
+    segments: List[TranscriptSegment] = []
+    current: Optional[TranscriptSegment] = None
+
+    for w in words:
+        text = w.get("word", "").strip()
+        if not text:
+            continue
+
+        w_start = float(w["start"])
+        w_end = float(w["end"])
+
+        speaker_raw = assign_speaker_to_word(w_start, diarization_segments)
+        speaker = speaker_map.get(speaker_raw, speaker_raw)
+        role = roles.get(speaker, "KH")
+
+        if current is None:
+            current = TranscriptSegment(
+                start=w_start,
+                end=w_end,
+                speaker=speaker,
+                role=role,
+                text=text,
+            )
+            continue
+
+        gap = w_start - current.end
+
+        if speaker == current.speaker and gap <= max_word_gap_s:
+            current.text += " " + text
+            current.end = max(current.end, w_end)
+        else:
+            segments.append(current)
+            current = TranscriptSegment(
+                start=w_start,
+                end=w_end,
+                speaker=speaker,
+                role=role,
+                text=text,
+            )
 
-    return {
-        spk: ("NV" if spk == agent else "KH")
-        for spk in dur
-    }
+    if current:
+        segments.append(current)
 
+    return segments
 
 def convert_audio_to_wav(audio_path: Path) -> Path:
     """Convert any audio to WAV 16kHz Mono using ffmpeg."""
@@ -89,14 +126,59 @@ def convert_audio_to_wav(audio_path: Path) -> Path:
     except subprocess.CalledProcessError as e:
         logger.error(f"FFmpeg conversion failed: {e}")
         return audio_path
-    
-    
+
+
 def format_timestamp(seconds: float) -> str:
     m = int(seconds // 60)
     s = seconds % 60
     return f"{m:02d}:{s:06.3f}"
 
 
+def pad_and_refine_tensor(
+    waveform: torch.Tensor,
+    sr: int,
+    start_s: float,
+    end_s: float,
+    pad_ms: int = 200,
+    silence_db_delta: float = 16,
+    min_duration_ms: int = 150,
+) -> Optional[Tuple[int, int]]:
+    """
+    Refine segment using energy on TORCH tensor.
+    Returns sample indices or None.
+    """
+    total_len = waveform.shape[1]
+
+    start_idx = max(int((start_s - pad_ms / 1000) * sr), 0)
+    end_idx = min(int((end_s + pad_ms / 1000) * sr), total_len)
+
+    if end_idx <= start_idx:
+        return None
+
+    segment  = waveform[:, start_idx:end_idx]
+    if segment .numel() == 0:
+        return None
+
+    # RMS energy
+    rms = torch.sqrt(torch.mean(segment ** 2) + 1e-9)
+    threshold = rms / silence_db_delta
+
+    energy = torch.abs(segment)
+    valid = torch.where(energy > threshold)[0]
+
+    if valid.numel() == 0:
+        return None
+
+
+    new_start = start_idx + valid[0].item()
+    new_end = start_idx + valid[-1].item()
+
+    if new_end - new_start < int(min_duration_ms / 1000 * sr):
+        return None
+
+    return new_start, new_end
+
+
 # =========================
 # Processor
 # =========================
@@ -105,9 +187,20 @@ class Processor:
     async def process_audio(
         cls,
         audio_path: Path,
+        model_name: str = "PhoWhisper VI Finetuned",
         language: str = "vi",
+ 
+        # VAD options
+        vad_filter: bool = True,
+        vad_min_silence_ms: int = 1000,
+        vad_speech_pad_ms: int = 400,
+        vad_min_speech_ms: int = 250,
+        vad_threshold: float = 0.5,
+        # Generation options
         beam_size: int = 5,
         temperature: float = 0.0,
+        best_of: int = 5,
+        initial_prompt: Optional[str] = None,
     ) -> ProcessingResult:
 
         import asyncio
@@ -120,97 +213,172 @@ class Processor:
 
         # 2: Load audio
         y, sr = librosa.load(wav_path, sr=16000, mono=True)
-        waveform = torch.from_numpy(y).unsqueeze(0)
+        if y.size == 0:
+            raise ValueError("Empty audio")
+        waveform = torch.from_numpy(y).unsqueeze(0).float()
         duration = len(y) / sr
-  
+        
         # 3: Diarization
+        logger.info("Step 3: Running diarization...")
+        
         diarization: DiarizationResult = await DiarizationService.diarize_async(wav_path)
-        diarization_segments = diarization.segments or [
-            SpeakerSegment(0.0, duration, "SPEAKER_0")
-        ]
-
-        diarization_segments.sort(key=lambda s: s.start)
 
-        diarization_segments = [
-            SpeakerSegment(
-                *pad_and_refine_tensor(waveform, sr, s.start, s.end),
-                speaker=s.speaker,
-            )
-            for s in diarization_segments
-        ]
+        diarization_segments = diarization.segments or []
+        speakers = diarization.speakers or []
+        roles = diarization.roles or {}
+
+        if not diarization_segments:
+            diarization_segments = [SpeakerSegment(0.0, duration, "SPEAKER_0")]
+            speakers = ["SPEAKER_0"]
+            roles = {"SPEAKER_0": "KH"}
+
+        diarization_segments.sort(key=lambda x: x.start)
+
+        # 4: Refine segment boundaries
+        refined_segments: List[SpeakerSegment] = []
+
+        for seg in diarization_segments:
+            refined = pad_and_refine_tensor(waveform, sr, seg.start, seg.end)
+
+            if refined:
+                s, e = refined
+                if e > s:
+                    refined_segments.append(
+                        SpeakerSegment(
+                            start=s / sr,
+                            end=e / sr,
+                            speaker=seg.speaker,
+                        )
+                    )
+                else:
+                    refined_segments.append(seg)
+            else:
+                refined_segments.append(seg)
 
-        # 4. SPEAKER MAP 
-        raw_speakers = sorted({s.speaker for s in diarization_segments})
-        speaker_map = {s: f"Speaker {i+1}" for i, s in enumerate(raw_speakers)}
-        
-        diarization_segments = [
-            SpeakerSegment(
-                start=s.start,
-                end=s.end,
-                speaker=speaker_map[s.speaker]
-            )
-            for s in diarization_segments
-        ]
-        # 5. Roles infer
-        roles = infer_roles_from_diarization(diarization_segments)
-
-        result = await TranscriptionService.transcribe_with_words_async(
-            audio_array=y,
-            language=language,
-            beam_size=beam_size,
-            temperature=temperature
-        )
-
-        words: List[WordTimestamp] = [
-            WordTimestamp(
-                word=w["word"],
-                start=float(w["start"]),
-                end=float(w["end"]),
-            )
-            for w in result.get("words", [])
-            if w.get("word")
-        ]
+        if not refined_segments:
+            refined_segments = diarization_segments
 
+    
+        # 5. Normalize speakers
+        raw_speakers = sorted({seg.speaker for seg in refined_segments})
 
-        aligned_segments = AlignmentService.align_precision(
-            words,
-            diarization_segments
-        )
+        speaker_map = {
+            spk: f"Speaker {i+1}"
+            for i, spk in enumerate(raw_speakers)
+        }
 
-        segments = [
-            TranscriptSegment(
-                start=s.start,
-                end=s.end,
-                speaker=s.speaker,
-                role=roles.get(s.speaker, "KH"),
-                text=s.text,
+        speakers = list(speaker_map.values())
+
+        # 6. NORMALIZE ROLES
+        speaker_duration = defaultdict(float)
+        for seg in refined_segments:
+            speaker_duration[seg.speaker] += seg.end - seg.start
+
+        logger.info(f"speaker_duration(raw) = {speaker_duration}")
+
+        if speaker_duration:
+            agent_raw = max(speaker_duration, key=speaker_duration.get)
+
+            roles = {
+                speaker_map[spk]: ("NV" if spk == agent_raw else "KH")
+                for spk in speaker_duration
+            }
+        else:
+            roles = {}
+
+        # Default fallback
+        for label in speakers:
+            roles.setdefault(label, "KH")
+
+        logger.info(f"roles(mapped) = {roles}")
+
+
+            
+        # 7: Transcribe
+        vad_options = None
+        if vad_filter:
+            vad_options = {
+                "min_silence_duration_ms": vad_min_silence_ms,
+                "speech_pad_ms": vad_speech_pad_ms,
+                "min_speech_duration_ms": vad_min_speech_ms,
+                "threshold": vad_threshold
+            }
+
+        processed_segments: List[TranscriptSegment] = []
+
+        for seg in refined_segments:
+            start = int(seg.start * sr)
+            end = int(seg.end * sr)
+
+            if end <= start:
+                continue
+
+            audio_slice = y[start:end]
+            if audio_slice.size < sr * 0.25:
+                continue
+
+            try:
+                text = await TranscriptionService.transcribe_with_words_async(
+                    audio_array=audio_slice,
+                    model_name=model_name,
+                    language=language,
+                    vad_options=vad_options,
+                    beam_size=beam_size,
+                    temperature=temperature,
+                    best_of=best_of,
+                    initial_prompt=initial_prompt,
+                )
+            except Exception as e:
+                logger.error(f"Transcribe error: {e}")
+                continue
+
+            if not text or not text.strip():
+                continue
+            
+            label = speaker_map.get(seg.speaker, seg.speaker)
+
+            processed_segments.append(
+                TranscriptSegment(
+                    start=seg.start,
+                    end=seg.end,
+                    speaker=label,
+                    role=roles[label],
+                    text=text.strip(),
+                )
             )
-            for s in aligned_segments
-        ]
-
-        segments = cls._filter_segments_with_context(segments)
 
-        segments = cls._merge_adjacent_segments(
-            segments,
+        if not processed_segments:
+            processed_segments = [
+                TranscriptSegment(
+                    start=0.0,
+                    end=duration,
+                    speaker=speakers[0],
+                    role=roles[speakers[0]],
+                    text="(No speech detected)"
+                )
+            ]
+        
+        processed_segments = cls._merge_adjacent_segments(
+            processed_segments,
             max_gap_s=0.6
         )
 
-        processing_time = time.time() - t0      
-        speakers = sorted({s.speaker for s in segments})
+        processed_segments = cls._filter_segments_with_context(processed_segments)
+        processing_time = time.time() - t0
         
         txt_content = cls._generate_txt(
-            segments,
-            len(speaker_map),
+            processed_segments,
+            len(speakers),
             processing_time,
             duration,
             roles
         )
 
-        csv_content = cls._generate_csv(segments)
+        csv_content = cls._generate_csv(processed_segments)
 
         return ProcessingResult(
-            segments=segments,
-            speaker_count=len(speaker_map),
+            segments=processed_segments,
+            speaker_count=len(speakers),
             duration=duration,
             processing_time=processing_time,
             speakers=speakers,
@@ -224,36 +392,44 @@ class Processor:
     def _is_meaningful_segment(
         seg: TranscriptSegment,
         min_duration_s: float = 0.6,
-        min_words: int = 3,
+        eps: float = 0.05,
+        min_words: int = 3
     ) -> bool:
-        if seg.end - seg.start >= min_duration_s:
+        duration = seg.end - seg.start
+        word_count = len(seg.text.split())
+
+        if duration + eps >= min_duration_s:
             return True
-        if len(seg.text.split()) >= min_words:
+        if word_count >= min_words:
             return True
         if seg.role == "KH":
             return True
+
         return False
 
     @classmethod
     def _filter_segments_with_context(
         cls,
-        segments: List[TranscriptSegment],
+        segments: List[TranscriptSegment]
     ) -> List[TranscriptSegment]:
-
         if not segments:
             return segments
-
+        
         segments = sorted(segments, key=lambda s: s.start)
         result = []
+        n = len(segments)
 
         for i, seg in enumerate(segments):
-            prev = segments[i - 1] if i > 0 else None
-            next_ = segments[i + 1] if i < len(segments) - 1 else None
+            prev_seg = segments[i - 1] if i > 0 else None
+            next_seg = segments[i + 1] if i < n - 1 else None
 
             if cls._is_meaningful_segment(seg):
                 result.append(seg)
-            elif prev and next_ and prev.speaker == seg.speaker == next_.speaker:
-                result.append(seg)
+                continue
+
+            if prev_seg and next_seg:
+                if prev_seg.speaker == seg.speaker == next_seg.speaker:
+                    result.append(seg)
 
         return result
     
@@ -261,20 +437,30 @@ class Processor:
     @staticmethod
     def _merge_adjacent_segments(
         segments: List[TranscriptSegment],
-        max_gap_s: float = 0.5,
+        max_gap_s: float = 0.5
     ) -> List[TranscriptSegment]:
-
+        """
+        Merge adjacent segments if:
+        - same speaker
+        - gap <= max_gap_s
+        """
         if not segments:
             return segments
 
+        segments = sorted(segments, key=lambda s: s.start)
         merged = [segments[0]]
 
         for seg in segments[1:]:
             prev = merged[-1]
+
             gap = seg.start - prev.end
 
-            if seg.speaker == prev.speaker and gap <= max_gap_s:
-                prev.text += " " + seg.text
+            if (
+                seg.speaker == prev.speaker
+                and gap <= max_gap_s
+            ):
+                # MERGE
+                prev.text = f"{prev.text} {seg.text}".strip()
                 prev.end = max(prev.end, seg.end)
             else:
                 merged.append(seg)

app/services/transcription.py

@@ -3,14 +3,11 @@ Transcription service using faster-whisper.
 Supports multiple Vietnamese Whisper models with caching.
 """
 import logging
-import torch
 from typing import Dict, Optional, List
 from dataclasses import dataclass
 
 import numpy as np
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
-from peft import PeftModel
-
+from faster_whisper import WhisperModel
 
 from app.core.config import get_settings
 
@@ -20,9 +17,7 @@ settings = get_settings()
 
 # Available Whisper models for Vietnamese
 AVAILABLE_MODELS = {
-    
-    "Whisper-LoRA": settings.whisper_lora_model_dir
-    
+    "PhoWhisper VI Finetuned": settings.default_whisper_model
 }
 
 
@@ -40,88 +35,134 @@ class TranscriptionService:
     Supports multiple models with caching.
     """
     
-    _model = None
-    _processor = None
-    _device = "cuda" if torch.cuda.is_available() else "cpu"
+    _models: Dict[str, WhisperModel] = {}
     
     @classmethod
-    def get_model(cls):
-        if cls._model is not None:
-            return cls._model, cls._processor
-
-        model_dir = AVAILABLE_MODELS["Whisper-LoRA"]
-
-        logger.info(f"Loading Whisper + LoRA from {model_dir}")
-        logger.info(f"Device: {cls._device}")
-
-        base_model = WhisperForConditionalGeneration.from_pretrained(model_dir)
-        model = PeftModel.from_pretrained(base_model, model_dir)
-
-        model.to(cls._device)
-        model.eval()
-
-        processor = WhisperProcessor.from_pretrained(model_dir)
-
-        cls._model = model
-        cls._processor = processor
-
-        logger.info("Whisper + LoRA loaded successfully")
-        return model, processor
+    def get_model(cls, model_name: str = None) -> WhisperModel:
+        """
+        Get or load a Whisper model (lazy loading with caching).
+        
+        Args:
+            model_name: Name of the model from AVAILABLE_MODELS
+            
+        Returns:
+            Loaded WhisperModel instance
+        """
+        
+        if model_name is None:
+            model_name = settings.default_whisper_model
+            
+        cache_key = f"{model_name}_{settings.resolved_compute_type}"
+        
+        if cache_key in cls._models:
+            return cls._models[cache_key]
+        
+        # Get model path
+        if model_name in AVAILABLE_MODELS:
+            model_path = AVAILABLE_MODELS[model_name]
+        else:
+            # Fallback to first available model
+            model_name = list(AVAILABLE_MODELS.keys())[0]
+            model_path = AVAILABLE_MODELS[model_name]
+        
+        logger.info(f"Loading Whisper model: {model_name} ({model_path})")
+        logger.debug(f"Device: {settings.resolved_device}, Compute type: {settings.resolved_compute_type}")
+        
+        model = WhisperModel(
+            model_path,
+            device=settings.resolved_device,
+            compute_type=settings.resolved_compute_type,
+        )
+        
+        cls._models[cache_key] = model
+        logger.info(f"Whisper model loaded: {model_name}")
+        
+        return model
     
     @classmethod
-    def is_loaded(cls) -> bool:
-        return cls._model is not None
+    def is_loaded(cls, model_name: str = None) -> bool:
+        if model_name is None:
+            model_name = settings.default_whisper_model
+        """Check if a model is loaded."""
+        cache_key = f"{model_name}_{settings.resolved_compute_type}"
+        return cache_key in cls._models
     
     @classmethod
-    def preload_model(cls) -> None:
-        cls.get_model()
+    def preload_model(cls, model_name: str = None) -> None:
+        """Preload a model during startup."""
+        if model_name is None:
+            model_name = settings.default_whisper_model
+        try:
+            cls.get_model(model_name)
+        except Exception as e:
+            logger.error(f"Failed to preload Whisper model: {e}")
+            raise
     
     @classmethod
     def transcribe_with_words(
         cls,
         audio_array: np.ndarray,
+        model_name: str = None,
         language: str = "vi",
+        vad_options: Optional[dict] = None,
         beam_size: int = 5,
         temperature: float = 0.0,
+        best_of: int = 5,
+        initial_prompt: Optional[str] = None,
     ) -> Dict:
-        model, processor = cls.get_model()
+        """
+        Transcribe audio and return word-level timestamps.
+        """
+        model = cls.get_model(model_name)
 
-        if audio_array.ndim > 1:
-            audio_array = np.mean(audio_array, axis=0)
+        vad_filter = vad_options if vad_options else False
+        prompt = initial_prompt.strip() if initial_prompt and initial_prompt.strip() else None
 
-        inputs = processor(
+        segments_gen, info = model.transcribe(
             audio_array,
-            sampling_rate=16000,
-            return_tensors="pt"
-        ).input_features.to(cls._device)
-
-        forced_decoder_ids = processor.get_decoder_prompt_ids(
-            language=language,
-            task="transcribe"
+            language=language if language != "auto" else None,
+            beam_size=beam_size,
+            temperature=temperature,
+            best_of=best_of,
+
+            # QA / Stability
+            condition_on_previous_text=False,
+            no_speech_threshold=0.6,
+
+            word_timestamps=True,
+
+            # VAD
+            vad_filter=vad_filter,
+            vad_parameters=dict(
+                threshold=settings.vad_threshold,
+                min_speech_duration_ms=settings.vad_min_speech_duration_ms,
+                min_silence_duration_ms=settings.vad_min_silence_duration_ms,
+            ),
+
+            initial_prompt=prompt,
         )
 
-        with torch.no_grad():
-            generated_ids = model.generate(
-                inputs,
-                forced_decoder_ids=forced_decoder_ids,
-                num_beams=beam_size,
-                temperature=temperature,
-                max_new_tokens=settings.whisper_max_new_tokens,
-            )
+        words = []
+        full_text = []
+
+        for seg in segments_gen:
+            if seg.text:
+                full_text.append(seg.text.strip())
 
-        text = processor.batch_decode(
-            generated_ids,
-            skip_special_tokens=True
-        )[0].strip()
+            if hasattr(seg, "words") and seg.words:
+                for w in seg.words:
+                    if not w.word.strip():
+                        continue
+                    words.append({
+                        "word": w.word.strip(),
+                        "start": float(w.start),
+                        "end": float(w.end),
+                    })
 
         return {
-            "text": text,
-            "words": [],
-            "info": {
-                "engine": "transformers-whisper-lora",
-                "language": language,
-                "beam_size": beam_size,
-            },
+            "text": " ".join(full_text).strip(),
+            "words": words,
+            "info": info,
         }
 
     
@@ -129,15 +170,35 @@ class TranscriptionService:
     async def transcribe_with_words_async(
         cls,
         audio_array: np.ndarray,
-        **kwargs
-    ) -> Dict:
+        model_name: str = None,
+        language: str = "vi",
+        vad_options: Optional[dict] = None,
+        beam_size: int = 5,
+        temperature: float = 0.0,
+        best_of: int = 5,
+        initial_prompt: Optional[str] = None,
+    ) -> str:
+        """
+        Async wrapper for transcription (runs in thread pool).
+        """
         import asyncio
+        
         loop = asyncio.get_event_loop()
         return await loop.run_in_executor(
             None,
-            lambda: cls.transcribe_with_words(audio_array, **kwargs)
+            lambda: cls.transcribe_with_words(
+                audio_array,
+                model_name=model_name,
+                language=language,
+                vad_options=vad_options,
+                beam_size=beam_size,
+                temperature=temperature,
+                best_of=best_of,
+                initial_prompt=initial_prompt
+            )
         )
-
+    
     @classmethod
     def get_available_models(cls) -> Dict[str, str]:
+        """Return list of available models."""
         return AVAILABLE_MODELS.copy()