Pooya-Fallah/FastConformer_Persian_Streaming

1import torch2import nemo.collections.asr as nemo_asr3import gc4import numpy as np5import torchaudio6import gradio as gr7 8pretrained_model_path="./stt_fa_fastconformer_hybrid_large_finetuned.nemo"9 10# Clear up memory11torch.cuda.empty_cache()12gc.collect()13model = nemo_asr.models.EncDecHybridRNNTCTCModel.restore_from(pretrained_model_path)14device = 'cuda' if torch.cuda.is_available() else 'cpu'15# device = 'cpu'  # You can transcribe even longer samples on the CPU, though it will take much longer !16model = model.to(device)17model.freeze()18 19def transcribe(stream, new_chunk):20    if new_chunk is None:21        return None, ""22    # 'audio' is a tuple: (sample_rate, data)23    sample_rate, data = new_chunk24 25    # Ensure the model is on the correct device26    device = 'cuda' if torch.cuda.is_available() else 'cpu'27 28    # Convert audio data to the expected format29    if isinstance(data, np.ndarray):30        audio_tensor = torch.tensor(data, dtype=torch.float32)31    else:32        raise ValueError("Audio data must be a numpy array")33 34    # Resample if sample rate is not 1600035    target_sample_rate = 1600036    if sample_rate != target_sample_rate:37        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate)38        audio_tensor = resampler(audio_tensor)39 40    if stream is not None:41        stream['audio'] = torch.cat([stream['audio'], audio_tensor], dim=-1)42    else:43        stream = {"text": ""}44        stream['audio'] = audio_tensor45 46 47    max_length = 5 * target_sample_rate  # 5 seconds48    new_text = ""49 50    # Process all chunks that fit max_length51    while stream['audio'].shape[-1] > max_length:52        # Extract first max_length samples53        audio_chunk = stream['audio'][..., :max_length]54 55        # Transcribe56        with torch.no_grad():57            transcript = model.transcribe(audio_chunk)  # Add batch dimension if needed58 59        # Update text (adjust based on model's output format)60        new_text += " " + transcript[0][0].strip()  # Example adjustment61 62        # Remove processed audio from buffer63        stream['audio'] = stream['audio'][..., max_length:]64 65    stream['text'] += new_text66    return stream, stream['text'].strip()67 68 69interface = gr.Interface(70    fn=transcribe,71    inputs=['state', gr.Audio(sources="microphone", streaming=True, type="numpy")],72    outputs=["state", "text"],73    live=True,74)75 76interface.launch()