Upload 2 files

app.py

@@ -0,0 +1,267 @@
+# app.py
+
+# -*- coding: utf-8 -*-
+"""
+Vietnamese End-to-End Speech Recognition using Wav2Vec 2.0 with Speaker Diarization.
+Streamlit Application with merged speaker segments and timestamps.
+"""
+
+import os
+import zipfile
+import torch
+import soundfile as sf
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
+import kenlm
+from pyctcdecode import Alphabet, BeamSearchDecoderCTC, LanguageModel
+from huggingface_hub import hf_hub_download
+import streamlit as st
+import numpy as np
+import librosa
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+@st.cache_resource(show_spinner=False)
+def load_model_and_tokenizer(cache_dir='./cache/'):
+    st.info("Loading processor and model...")
+    processor = Wav2Vec2Processor.from_pretrained(
+        "nguyenvulebinh/wav2vec2-base-vietnamese-250h",
+        cache_dir=cache_dir
+    )
+    model = Wav2Vec2ForCTC.from_pretrained(
+        "nguyenvulebinh/wav2vec2-base-vietnamese-250h",
+        cache_dir=cache_dir
+    )
+
+    st.info("Downloading language model...")
+    lm_zip_file = hf_hub_download(
+        repo_id="nguyenvulebinh/wav2vec2-base-vietnamese-250h",
+        filename="vi_lm_4grams.bin.zip",
+        cache_dir=cache_dir
+    )
+
+    st.info("Extracting language model...")
+    with zipfile.ZipFile(lm_zip_file, 'r') as zip_ref:
+        zip_ref.extractall(cache_dir)
+
+    lm_file = os.path.join(cache_dir, 'vi_lm_4grams.bin')
+    if not os.path.isfile(lm_file):
+        raise FileNotFoundError(f"Language model file not found: {lm_file}")
+
+    st.success("Processor, model, and language model loaded successfully.")
+    return processor, model, lm_file
+
+@st.cache_resource(show_spinner=False)
+def get_decoder_ngram_model(_tokenizer, ngram_lm_path):
+    st.info("Building decoder with n-gram language model...")
+    vocab_dict = _tokenizer.get_vocab()
+    sorted_vocab = sorted((value, key) for (key, value) in vocab_dict.items())
+    vocab_list = [token for _, token in sorted_vocab][:-2]  # Exclude special tokens
+
+    alphabet = Alphabet.build_alphabet(vocab_list)
+    lm_model = kenlm.Model(ngram_lm_path)
+    decoder = BeamSearchDecoderCTC(alphabet, language_model=LanguageModel(lm_model))
+    st.success("Decoder built successfully.")
+    return decoder
+
+def transcribe_chunk(model, processor, decoder, speech_chunk, sampling_rate):
+    if speech_chunk.ndim > 1:
+        speech_chunk = np.mean(speech_chunk, axis=1)
+    speech_chunk = speech_chunk.astype(np.float32)
+
+    target_sr = 16000
+    if sampling_rate != target_sr:
+        speech_chunk = librosa.resample(speech_chunk, orig_sr=sampling_rate, target_sr=target_sr)
+        sampling_rate = target_sr
+
+    MIN_DURATION = 0.5  # seconds
+    MIN_SAMPLES = int(MIN_DURATION * sampling_rate)
+
+    if len(speech_chunk) < MIN_SAMPLES:
+        # Pad with zeros
+        padding = MIN_SAMPLES - len(speech_chunk)
+        speech_chunk = np.pad(speech_chunk, (0, padding), 'constant')
+
+    input_values = processor(
+        speech_chunk, sampling_rate=sampling_rate, return_tensors="pt"
+    ).input_values
+
+    with torch.no_grad():
+        logits = model(input_values).logits[0]
+
+    beam_search_output = decoder.decode(
+        logits.cpu().detach().numpy(),
+        beam_width=500
+    )
+    return beam_search_output
+
+def alternative_speaker_diarization(audio_file, num_speakers=2):
+    try:
+        # Use librosa to load the audio file
+        y, sr = librosa.load(audio_file, sr=None)
+
+        # Rough segmentation based on energy
+        intervals = librosa.effects.split(y, top_db=30)  # Adjust top_db as needed
+
+        # Merge very short intervals
+        MIN_INTERVAL_DURATION = 0.5  # seconds
+        MIN_SAMPLES = int(MIN_INTERVAL_DURATION * sr)
+        merged_intervals = []
+        for interval in intervals:
+            if merged_intervals and (interval[0] - merged_intervals[-1][1]) < MIN_SAMPLES:
+                merged_intervals[-1][1] = interval[1]
+            else:
+                merged_intervals.append([interval[0], interval[1]])
+
+        # Assign speakers cyclically
+        segments = []
+        for i, (start, end) in enumerate(merged_intervals):
+            speaker_id = i % num_speakers
+            start_time = start / sr
+            end_time = end / sr
+            segments.append((start_time, end_time, speaker_id))
+
+        return segments
+
+    except Exception as e:
+        st.error(f"Speaker diarization failed: {e}")
+        # Fallback to a simple equal-length segmentation
+        audio, sr = sf.read(audio_file)
+        total_duration = len(audio) / sr
+        segment_duration = total_duration / num_speakers
+
+        segments = []
+        for i in range(num_speakers):
+            start = i * segment_duration
+            end = (i + 1) * segment_duration
+            segments.append((start, end, i))
+
+        return segments
+
+def process_segments(audio_file, segments, model, processor, decoder, sampling_rate=16000):
+    speech, sr = sf.read(audio_file)
+    final_transcriptions = []
+
+    # Remove duplicate or overlapping segments
+    unique_segments = []
+    for segment in sorted(segments, key=lambda x: x[0]):
+        if not unique_segments or segment[0] >= unique_segments[-1][1]:
+            unique_segments.append(segment)
+
+    for start, end, speaker_id in unique_segments:
+        start_sample = int(start * sr)
+        end_sample = int(end * sr)
+        speech_chunk = speech[start_sample:end_sample]
+        transcript = transcribe_chunk(model, processor, decoder, speech_chunk, sr)
+
+        # Only add non-empty transcripts
+        if transcript.strip():
+            # Lưu (start, end, speaker_id, transcript)
+            final_transcriptions.append((start, end, speaker_id, transcript))
+
+    return final_transcriptions
+
+def format_timestamp(seconds):
+    # Định dạng thời gian thành MM:SS
+    total_seconds = int(seconds)
+    mm = total_seconds // 60
+    ss = total_seconds % 60
+    return f"{mm:02d}:{ss:02d}"
+
+def merge_speaker_segments(final_transcriptions):
+    # Gộp các đoạn cùng speaker liên tiếp
+    if not final_transcriptions:
+        return []
+
+    merged_results = []
+    prev_start, prev_end, prev_speaker_id, prev_text = final_transcriptions[0]
+
+    for i in range(1, len(final_transcriptions)):
+        start, end, speaker_id, text = final_transcriptions[i]
+        if speaker_id == prev_speaker_id:
+            # Cùng speaker, gộp đoạn
+            prev_end = end
+            prev_text += " " + text
+        else:
+            # Khác speaker
+            merged_results.append((prev_start, prev_end, prev_speaker_id, prev_text))
+            prev_start, prev_end, prev_speaker_id, prev_text = start, end, speaker_id, text
+
+    # Thêm đoạn cuối cùng
+    merged_results.append((prev_start, prev_end, prev_speaker_id, prev_text))
+
+    return merged_results
+
+def main():
+    st.title("🇻🇳 Vietnamese Speech Recognition with Speaker Diarization (with merging & timestamps)")
+
+    st.write("""
+    Upload an audio file, select the number of speakers, and get the transcribed text with timestamps and merged segments for each speaker.
+    """)
+
+    # Sidebar for inputs
+    st.sidebar.header("Input Parameters")
+    uploaded_file = st.sidebar.file_uploader("Upload Audio File", type=["wav", "mp3", "flac", "m4a"])
+    num_speakers = st.sidebar.slider("Number of Speakers", min_value=1, max_value=5, value=2, step=1)
+
+    if uploaded_file is not None:
+        # Save the uploaded file to a temporary location
+        temp_audio_path = "temp_audio_file"
+        with open(temp_audio_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+
+        # Display audio player
+        st.audio(uploaded_file, format='audio/wav')
+
+        if st.button("Transcribe"):
+            with st.spinner("Processing..."):
+                try:
+                    # Load models
+                    processor, model, lm_file = load_model_and_tokenizer()
+                    decoder = get_decoder_ngram_model(processor.tokenizer, lm_file)
+
+                    # Speaker diarization
+                    segments = alternative_speaker_diarization(temp_audio_path, num_speakers=num_speakers)
+
+                    if not segments:
+                        st.warning("No speech segments detected.")
+                        return
+
+                    # Process segments
+                    final_transcriptions = process_segments(temp_audio_path, segments, model, processor, decoder)
+
+                    # Merge consecutive segments of the same speaker
+                    merged_results = merge_speaker_segments(final_transcriptions)
+
+                    # Display results
+                    if merged_results:
+                        st.success("Transcription Completed!")
+                        transcription_text = ""
+                        for start_time, end_time, speaker_id, transcript in merged_results:
+                            start_str = format_timestamp(start_time)
+                            end_str = format_timestamp(end_time)
+                            line = f"{start_str} - {end_str} - Speaker {speaker_id + 1}: {transcript}"
+                            st.markdown(line)
+                            transcription_text += line + "\n"
+
+                        # Provide download link
+                        st.download_button(
+                            label="Download Transcription",
+                            data=transcription_text,
+                            file_name="transcription.txt",
+                            mime="text/plain"
+                        )
+                    else:
+                        st.warning("No transcriptions available.")
+
+                except Exception as e:
+                    st.error(f"An error occurred during processing: {e}")
+
+            # Optionally, remove the temporary file after processing
+            if os.path.exists(temp_audio_path):
+                os.remove(temp_audio_path)
+    else:
+        st.info("Please upload an audio file to get started.")
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,311 @@
+aiohappyeyeballs==2.4.4
+aiohttp==3.11.10
+aioice==0.9.0
+aiortc==1.9.0
+aiosignal==1.3.1
+alembic==1.14.0
+altair==5.5.0
+antlr4-python3-runtime==4.9.3
+asteroid-filterbanks==0.4.0
+attrs==24.2.0
+audioread==3.0.1
+av==12.3.0
+blinker==1.9.0
+cachetools==5.5.0
+certifi==2024.8.30
+cffi==1.17.1
+charset-normalizer==3.4.0
+click==8.1.7
+colorlog==6.9.0
+contourpy==1.3.1
+coverage==5.5
+cryptography==44.0.0
+cycler==0.12.1
+datasets==3.1.0
+decorator==5.1.1
+deprecation==2.1.0
+dill==0.3.8
+dnspython==2.7.0
+docopt==0.6.2
+einops==0.8.0
+eyed3==0.9.7
+ffmpeg-python==0.2.0
+filelock==3.16.1
+filetype==1.2.0
+fonttools==4.55.2
+frozenlist==1.5.0
+fsspec==2024.9.0
+future==1.0.0
+gitdb==4.0.11
+GitPython==3.1.43
+google-crc32c==1.6.0
+hmmlearn==0.3.3
+huggingface-hub==0.26.3
+HyperPyYAML==1.2.2
+hypothesis==6.122.1
+idna==3.10
+ifaddr==0.2.0
+imbalanced-learn==0.12.4
+imblearn==0.0
+Jinja2==3.1.4
+joblib==1.4.2
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
+julius==0.2.7
+kenlm @ git+https://github.com/kpu/kenlm.git@f6c947dc943859e265fabce886232205d0fb2b37
+kiwisolver==1.4.7
+lazy_loader==0.4
+librosa==0.10.2.post1
+lightning==2.4.0
+lightning-utilities==0.11.9
+llvmlite==0.43.0
+Mako==1.3.8
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+matplotlib==3.9.3
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.1.0
+multidict==6.1.0
+multiprocess==0.70.16
+narwhals==1.15.2
+networkx==3.4.2
+numba==0.60.0
+numpy==1.26.4
+omegaconf==2.3.0
+optuna==4.1.0
+packaging==24.2
+pandas==2.2.3
+pillow==11.0.0
+platformdirs==4.3.6
+playsound==1.3.0
+plotly==5.24.1
+pooch==1.8.2
+primePy==1.3
+propcache==0.2.1
+protobuf==5.29.1
+pyannote.audio==3.3.2
+pyannote.core==5.0.0
+pyannote.database==5.1.0
+pyannote.metrics==3.2.1
+pyannote.pipeline==3.0.1
+pyarrow==18.1.0
+pyAudioAnalysis==0.3.14
+pycparser==2.22
+pyctcdecode==0.5.0
+pydeck==0.9.1
+pydub==0.25.1
+pyee==12.1.1
+Pygments==2.18.0
+pygtrie==2.5.0
+pylibsrtp==0.10.0
+pyobjc==10.3.2
+pyobjc-core==10.3.2
+pyobjc-framework-Accessibility==10.3.2
+pyobjc-framework-Accounts==10.3.2
+pyobjc-framework-AddressBook==10.3.2
+pyobjc-framework-AdServices==10.3.2
+pyobjc-framework-AdSupport==10.3.2
+pyobjc-framework-AppleScriptKit==10.3.2
+pyobjc-framework-AppleScriptObjC==10.3.2
+pyobjc-framework-ApplicationServices==10.3.2
+pyobjc-framework-AppTrackingTransparency==10.3.2
+pyobjc-framework-AudioVideoBridging==10.3.2
+pyobjc-framework-AuthenticationServices==10.3.2
+pyobjc-framework-AutomaticAssessmentConfiguration==10.3.2
+pyobjc-framework-Automator==10.3.2
+pyobjc-framework-AVFoundation==10.3.2
+pyobjc-framework-AVKit==10.3.2
+pyobjc-framework-AVRouting==10.3.2
+pyobjc-framework-BackgroundAssets==10.3.2
+pyobjc-framework-BrowserEngineKit==10.3.2
+pyobjc-framework-BusinessChat==10.3.2
+pyobjc-framework-CalendarStore==10.3.2
+pyobjc-framework-CallKit==10.3.2
+pyobjc-framework-CFNetwork==10.3.2
+pyobjc-framework-Cinematic==10.3.2
+pyobjc-framework-ClassKit==10.3.2
+pyobjc-framework-CloudKit==10.3.2
+pyobjc-framework-Cocoa==10.3.2
+pyobjc-framework-Collaboration==10.3.2
+pyobjc-framework-ColorSync==10.3.2
+pyobjc-framework-Contacts==10.3.2
+pyobjc-framework-ContactsUI==10.3.2
+pyobjc-framework-CoreAudio==10.3.2
+pyobjc-framework-CoreAudioKit==10.3.2
+pyobjc-framework-CoreBluetooth==10.3.2
+pyobjc-framework-CoreData==10.3.2
+pyobjc-framework-CoreHaptics==10.3.2
+pyobjc-framework-CoreLocation==10.3.2
+pyobjc-framework-CoreMedia==10.3.2
+pyobjc-framework-CoreMediaIO==10.3.2
+pyobjc-framework-CoreMIDI==10.3.2
+pyobjc-framework-CoreML==10.3.2
+pyobjc-framework-CoreMotion==10.3.2
+pyobjc-framework-CoreServices==10.3.2
+pyobjc-framework-CoreSpotlight==10.3.2
+pyobjc-framework-CoreText==10.3.2
+pyobjc-framework-CoreWLAN==10.3.2
+pyobjc-framework-CryptoTokenKit==10.3.2
+pyobjc-framework-DataDetection==10.3.2
+pyobjc-framework-DeviceCheck==10.3.2
+pyobjc-framework-DictionaryServices==10.3.2
+pyobjc-framework-DiscRecording==10.3.2
+pyobjc-framework-DiscRecordingUI==10.3.2
+pyobjc-framework-DiskArbitration==10.3.2
+pyobjc-framework-DVDPlayback==10.3.2
+pyobjc-framework-EventKit==10.3.2
+pyobjc-framework-ExceptionHandling==10.3.2
+pyobjc-framework-ExecutionPolicy==10.3.2
+pyobjc-framework-ExtensionKit==10.3.2
+pyobjc-framework-ExternalAccessory==10.3.2
+pyobjc-framework-FileProvider==10.3.2
+pyobjc-framework-FileProviderUI==10.3.2
+pyobjc-framework-FinderSync==10.3.2
+pyobjc-framework-FSEvents==10.3.2
+pyobjc-framework-GameCenter==10.3.2
+pyobjc-framework-GameController==10.3.2
+pyobjc-framework-GameKit==10.3.2
+pyobjc-framework-GameplayKit==10.3.2
+pyobjc-framework-HealthKit==10.3.2
+pyobjc-framework-ImageCaptureCore==10.3.2
+pyobjc-framework-InputMethodKit==10.3.2
+pyobjc-framework-InstallerPlugins==10.3.2
+pyobjc-framework-InstantMessage==10.3.2
+pyobjc-framework-Intents==10.3.2
+pyobjc-framework-IntentsUI==10.3.2
+pyobjc-framework-IOBluetooth==10.3.2
+pyobjc-framework-IOBluetoothUI==10.3.2
+pyobjc-framework-IOSurface==10.3.2
+pyobjc-framework-iTunesLibrary==10.3.2
+pyobjc-framework-KernelManagement==10.3.2
+pyobjc-framework-LatentSemanticMapping==10.3.2
+pyobjc-framework-LaunchServices==10.3.2
+pyobjc-framework-libdispatch==10.3.2
+pyobjc-framework-libxpc==10.3.2
+pyobjc-framework-LinkPresentation==10.3.2
+pyobjc-framework-LocalAuthentication==10.3.2
+pyobjc-framework-LocalAuthenticationEmbeddedUI==10.3.2
+pyobjc-framework-MailKit==10.3.2
+pyobjc-framework-MapKit==10.3.2
+pyobjc-framework-MediaAccessibility==10.3.2
+pyobjc-framework-MediaLibrary==10.3.2
+pyobjc-framework-MediaPlayer==10.3.2
+pyobjc-framework-MediaToolbox==10.3.2
+pyobjc-framework-Metal==10.3.2
+pyobjc-framework-MetalFX==10.3.2
+pyobjc-framework-MetalKit==10.3.2
+pyobjc-framework-MetalPerformanceShaders==10.3.2
+pyobjc-framework-MetalPerformanceShadersGraph==10.3.2
+pyobjc-framework-MetricKit==10.3.2
+pyobjc-framework-MLCompute==10.3.2
+pyobjc-framework-ModelIO==10.3.2
+pyobjc-framework-MultipeerConnectivity==10.3.2
+pyobjc-framework-NaturalLanguage==10.3.2
+pyobjc-framework-NetFS==10.3.2
+pyobjc-framework-Network==10.3.2
+pyobjc-framework-NetworkExtension==10.3.2
+pyobjc-framework-NotificationCenter==10.3.2
+pyobjc-framework-OpenDirectory==10.3.2
+pyobjc-framework-OSAKit==10.3.2
+pyobjc-framework-OSLog==10.3.2
+pyobjc-framework-PassKit==10.3.2
+pyobjc-framework-PencilKit==10.3.2
+pyobjc-framework-PHASE==10.3.2
+pyobjc-framework-Photos==10.3.2
+pyobjc-framework-PhotosUI==10.3.2
+pyobjc-framework-PreferencePanes==10.3.2
+pyobjc-framework-PushKit==10.3.2
+pyobjc-framework-Quartz==10.3.2
+pyobjc-framework-QuickLookThumbnailing==10.3.2
+pyobjc-framework-ReplayKit==10.3.2
+pyobjc-framework-SafariServices==10.3.2
+pyobjc-framework-SafetyKit==10.3.2
+pyobjc-framework-SceneKit==10.3.2
+pyobjc-framework-ScreenCaptureKit==10.3.2
+pyobjc-framework-ScreenSaver==10.3.2
+pyobjc-framework-ScreenTime==10.3.2
+pyobjc-framework-ScriptingBridge==10.3.2
+pyobjc-framework-SearchKit==10.3.2
+pyobjc-framework-Security==10.3.2
+pyobjc-framework-SecurityFoundation==10.3.2
+pyobjc-framework-SecurityInterface==10.3.2
+pyobjc-framework-SensitiveContentAnalysis==10.3.2
+pyobjc-framework-ServiceManagement==10.3.2
+pyobjc-framework-SharedWithYou==10.3.2
+pyobjc-framework-SharedWithYouCore==10.3.2
+pyobjc-framework-ShazamKit==10.3.2
+pyobjc-framework-Social==10.3.2
+pyobjc-framework-SoundAnalysis==10.3.2
+pyobjc-framework-Speech==10.3.2
+pyobjc-framework-SpriteKit==10.3.2
+pyobjc-framework-StoreKit==10.3.2
+pyobjc-framework-Symbols==10.3.2
+pyobjc-framework-SyncServices==10.3.2
+pyobjc-framework-SystemConfiguration==10.3.2
+pyobjc-framework-SystemExtensions==10.3.2
+pyobjc-framework-ThreadNetwork==10.3.2
+pyobjc-framework-UniformTypeIdentifiers==10.3.2
+pyobjc-framework-UserNotifications==10.3.2
+pyobjc-framework-UserNotificationsUI==10.3.2
+pyobjc-framework-VideoSubscriberAccount==10.3.2
+pyobjc-framework-VideoToolbox==10.3.2
+pyobjc-framework-Virtualization==10.3.2
+pyobjc-framework-Vision==10.3.2
+pyobjc-framework-WebKit==10.3.2
+pyOpenSSL==24.3.0
+pyparsing==3.2.0
+python-dateutil==2.9.0.post0
+pytorch-lightning==2.4.0
+pytorch-metric-learning==2.7.0
+pytz==2024.2
+PyYAML==6.0.2
+referencing==0.35.1
+regex==2024.11.6
+requests==2.32.3
+Resemblyzer==0.1.4
+rich==13.9.4
+rpds-py==0.22.3
+ruamel.yaml==0.18.6
+ruamel.yaml.clib==0.2.12
+safetensors==0.4.5
+scikit-learn==1.5.2
+scipy==1.14.1
+semver==3.0.2
+sentencepiece==0.2.0
+shellingham==1.5.4
+six==1.17.0
+smmap==5.0.1
+sortedcontainers==2.4.0
+sounddevice==0.5.1
+soundfile==0.12.1
+soxr==0.5.0.post1
+spectralcluster==0.2.22
+speechbrain==1.0.2
+SQLAlchemy==2.0.36
+streamlit==1.40.2
+streamlit-webrtc==0.47.9
+sympy==1.13.1
+tabulate==0.9.0
+tenacity==9.0.0
+tensorboardX==2.6.2.2
+threadpoolctl==3.5.0
+tokenizers==0.21.0
+toml==0.10.2
+torch==2.5.1
+torch-audiomentations==0.11.1
+torch_pitch_shift==1.2.5
+torchaudio==2.5.1
+torchmetrics==1.6.0
+tornado==6.4.2
+tqdm==4.67.1
+transformers==4.47.0
+typer==0.15.1
+typing==3.7.4.3
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+watchdog==6.0.0
+webrtcvad==2.0.10
+xxhash==3.5.0
+yarl==1.18.3