audio_processing.py

import gc
import torch
import torchaudio
import numpy as np
from transformers import (
    Wav2Vec2ForSequenceClassification, 
    AutoFeatureExtractor, 
    Wav2Vec2ForCTC, 
    AutoProcessor, 
    AutoTokenizer, 
    AutoModelForSeq2SeqLM
)
import spaces
import logging
from difflib import SequenceMatcher

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)



class AudioProcessor:
    def __init__(self, chunk_size=5, overlap=1, sample_rate=16000):
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.sample_rate = sample_rate
        self.previous_text = ""
        self.previous_lang = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def load_models(self):
        """Load all required models"""
        logger.info("Loading MMS models...")
        
        # Language identification model
        lid_processor = AutoFeatureExtractor.from_pretrained("facebook/mms-lid-256")
        lid_model = Wav2Vec2ForSequenceClassification.from_pretrained("facebook/mms-lid-256")
        
        # Transcription model
        mms_processor = AutoProcessor.from_pretrained("facebook/mms-1b-all")
        mms_model = Wav2Vec2ForCTC.from_pretrained("facebook/mms-1b-all")
        
        # Translation model
        translation_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
        translation_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
        
        return {
            'lid': (lid_model, lid_processor),
            'mms': (mms_model, mms_processor),
            'translation': (translation_model, translation_tokenizer)
        }

    @spaces.GPU(duration=60)
    def identify_language(self, audio_chunk, models):
        """Identify language of audio chunk"""
        lid_model, lid_processor = models['lid']
        inputs = lid_processor(audio_chunk, sampling_rate=16000, return_tensors="pt")
        lid_model.to(self.device)
        with torch.no_grad():
            outputs = lid_model(inputs.input_values.to(self.device)).logits
            lang_id = torch.argmax(outputs, dim=-1)[0].item()
            detected_lang = lid_model.config.id2label[lang_id]
        
        return detected_lang

    @spaces.GPU(duration=60)
    def transcribe_chunk(self, audio_chunk, language, models):
        """Transcribe audio chunk"""
        mms_model, mms_processor = models['mms']
        
        mms_processor.tokenizer.set_target_lang(language)
        mms_model.load_adapter(language)
        mms_model.to(self.device)
        inputs = mms_processor(audio_chunk, sampling_rate=16000, return_tensors="pt")
        
        with torch.no_grad():
            outputs = mms_model(inputs.input_values.to(self.device)).logits
            ids = torch.argmax(outputs, dim=-1)[0]
            transcription = mms_processor.decode(ids)
        
        return transcription

    @spaces.GPU(duration=60)
    def translate_text(self, text, models):
        """Translate text to English"""
        translation_model, translation_tokenizer = models['translation']
        
        inputs = translation_tokenizer(text, return_tensors="pt")
        inputs = inputs.to(self.device)
        translation_model.to(self.device)
        with torch.no_grad():
            outputs = translation_model.generate(
                **inputs,
                forced_bos_token_id=translation_tokenizer.convert_tokens_to_ids("eng_Latn"),
                max_length=100
            )
            translation = translation_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
        
        return translation

    def preprocess_audio(self, audio):
        """
        Create overlapping chunks with improved timing logic
        """
        chunk_samples = int(self.chunk_size * self.sample_rate)
        overlap_samples = int(self.overlap * self.sample_rate)

        chunks_with_times = []
        start_idx = 0

        while start_idx < len(audio):
            end_idx = min(start_idx + chunk_samples, len(audio))

            # Add padding for first chunk
            if start_idx == 0:
                chunk = audio[start_idx:end_idx]
                padding = torch.zeros(int(1 * self.sample_rate))
                chunk = torch.cat([padding, chunk])
            else:
                # Include overlap from previous chunk
                actual_start = max(0, start_idx - overlap_samples)
                chunk = audio[actual_start:end_idx]

            # Pad if necessary
            if len(chunk) < chunk_samples:
                chunk = torch.nn.functional.pad(chunk, (0, chunk_samples - len(chunk)))

            # Adjust time ranges to account for overlaps
            chunk_start_time = max(0, (start_idx / self.sample_rate) - self.overlap)
            chunk_end_time = min((end_idx / self.sample_rate) + self.overlap, len(audio) / self.sample_rate)

            chunks_with_times.append({
                'chunk': chunk,
                'start_time': start_idx / self.sample_rate,
                'end_time': end_idx / self.sample_rate,
                'transcribe_start': chunk_start_time,
                'transcribe_end': chunk_end_time
            })

            # Move to next chunk with smaller step size for better continuity
            start_idx += (chunk_samples - overlap_samples)

        return chunks_with_times

    
    @spaces.GPU(duration=60)
    def process_audio(self, audio_path, translate=False):
        """Main processing function"""
        try:
            # Load audio
            waveform, sample_rate = torchaudio.load(audio_path)
            if waveform.shape[0] > 1:
                waveform = torch.mean(waveform, dim=0)
            else: 
                waveform = waveform.squeeze(0)
                
            # Resample if necessary
            if sample_rate != self.sample_rate:
                resampler = torchaudio.transforms.Resample(
                    orig_freq=sample_rate,
                    new_freq=self.sample_rate
                )
                waveform = resampler(waveform)

            # if sample_rate != self.sample_rate:
            #     waveform = torchaudio.transforms.Resample(sample_rate, self.sample_rate)(waveform)

            # Load models
            models = self.load_models()
            
            # Process in chunks
            chunk_samples = int(self.chunk_size * self.sample_rate)
            overlap_samples = int(self.overlap * self.sample_rate)
            
            segments = []
            language_segments = []
            
            for i in range(0, len(waveform), chunk_samples - overlap_samples):
                chunk = waveform[i:i + chunk_samples]
                if len(chunk) < chunk_samples:
                    chunk = torch.nn.functional.pad(chunk, (0, chunk_samples - len(chunk)))
                
                # Process chunk
                start_time = i / self.sample_rate
                end_time = (i + len(chunk)) / self.sample_rate
                
                # Identify language
                language = self.identify_language(chunk, models)
                
                # Record language segment
                language_segments.append({
                    "language": language,
                    "start": start_time,
                    "end": end_time
                })
                
                # Transcribe
                transcription = self.transcribe_chunk(chunk, language, models)
                
                segment = {
                    "start": start_time,
                    "end": end_time,
                    "language": language,
                    "text": transcription,
                    "speaker": "Speaker"  # Simple speaker assignment
                }
                
                if translate:
                    translation = self.translate_text(transcription, models)
                    segment["translated"] = translation
                
                segments.append(segment)
                
                # Clean up GPU memory
                torch.cuda.empty_cache()
                gc.collect()
            
            # Merge nearby segments
            merged_segments = self.merge_segments(segments)
            
            return language_segments, merged_segments

        except Exception as e:
            logger.error(f"Error processing audio: {str(e)}")
            raise

    
    def merge_segments(self, segments, time_threshold=0.5, similarity_threshold=0.7):
        """Merge similar nearby segments"""
        if not segments:
            return segments
            
        merged = []
        current = segments[0]
        
        for next_segment in segments[1:]:
            if (next_segment['start'] - current['end'] <= time_threshold and 
                current['language'] == next_segment['language']):
                
                # Check text similarity
                matcher = SequenceMatcher(None, current['text'], next_segment['text'])
                similarity = matcher.ratio()
                
                if similarity > similarity_threshold:
                    # Merge segments
                    current['end'] = next_segment['end']
                    current['text'] = current['text'] + ' ' + next_segment['text']
                    if 'translated' in current and 'translated' in next_segment:
                        current['translated'] = current['translated'] + ' ' + next_segment['translated']
                else:
                    merged.append(current)
                    current = next_segment
            else:
                merged.append(current)
                current = next_segment
        
        merged.append(current)
        return merged