import os import json from pathlib import Path import whisperx import soundfile as sf import numpy as np import re import torch from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline import sys from dotenv import load_dotenv load_dotenv() token = os.getenv("HF_TOKEN") print("Start Preprocessing ... ...") sys.path.append('./CrisperWhisper/') from utils import adjust_pauses_for_hf_pipeline_output def generate_session_id(): session_root = "session_data" if not os.path.exists(session_root): os.makedirs(session_root) return "000001" existing_ids = [d for d in os.listdir(session_root) if os.path.isdir(os.path.join(session_root, d)) and d.isdigit()] if existing_ids: new_id = max(int(x) for x in existing_ids) + 1 else: new_id = 1 return f"{new_id:06d}" def assign_speakers(segments, diarization_segments): speaker_map = {} for segment in segments: segment_start = segment["start"] segment_end = segment["end"] max_overlap = 0 assigned_speaker = "Unknown" for _, diar in diarization_segments.iterrows(): speaker = diar["speaker"] diar_start = diar["start"] diar_end = diar["end"] overlap_start = max(segment_start, diar_start) overlap_end = min(segment_end, diar_end) overlap_duration = max(0, overlap_end - overlap_start) if overlap_duration > max_overlap: max_overlap = overlap_duration assigned_speaker = speaker speaker_map[segment_start] = assigned_speaker return speaker_map def load_audio_for_split(input_audio_file): if input_audio_file.lower().endswith('.mp3'): from pydub import AudioSegment audio_seg = AudioSegment.from_file(input_audio_file) sr = audio_seg.frame_rate samples = np.array(audio_seg.get_array_of_samples()).astype(np.float32) samples = samples / 32768.0 if audio_seg.channels > 1: samples = samples.reshape((-1, audio_seg.channels)) return samples, sr else: return sf.read(input_audio_file) def split_segment_by_sentences(segment): text = segment["text"] words = segment["words"] start_time = segment["start"] end_time = segment["end"] speaker = segment["speaker"] sentences = [s.strip() for s in text.split('.') if s.strip()] if len(sentences) <= 1: return [segment] new_segments = [] word_index = 0 for i, sentence in enumerate(sentences): if not sentence: continue sentence_words = [] sentence_text_clean = re.sub(r'[^\w\s]', '', sentence.lower()) sentence_word_tokens = sentence_text_clean.split() matched_words = 0 sentence_start = None sentence_end = None temp_word_index = word_index while temp_word_index < len(words) and matched_words < len(sentence_word_tokens): word_obj = words[temp_word_index] word_text_clean = re.sub(r'[^\w\s]', '', word_obj["word"].lower()) if word_text_clean == sentence_word_tokens[matched_words]: if sentence_start is None: sentence_start = word_obj["start"] sentence_end = word_obj["end"] sentence_words.append(word_obj) matched_words += 1 elif word_text_clean in sentence_word_tokens[matched_words:]: sentence_words.append(word_obj) if sentence_start is None: sentence_start = word_obj["start"] sentence_end = word_obj["end"] temp_word_index += 1 if sentence_start is None or sentence_end is None: total_duration = end_time - start_time sentence_duration = total_duration / len(sentences) sentence_start = start_time + i * sentence_duration sentence_end = start_time + (i + 1) * sentence_duration if i == len(sentences) - 1: sentence_end = end_time word_index = temp_word_index new_segment = { "start": round(sentence_start, 3), "end": round(sentence_end, 3), "speaker": speaker, "text": sentence + ".", "words": sentence_words } new_segments.append(new_segment) return new_segments def process_audio_file(input_audio_file, num_speakers, device="cuda"): print("Loading WhisperX model (English)...") model = whisperx.load_model("medium", device, language="en") audio = whisperx.load_audio(input_audio_file) print("Transcribing audio with WhisperX...") result = model.transcribe(audio) print("Performing forced alignment with WhisperX...") alignment_model, metadata = whisperx.load_align_model(language_code="en", device=device) result_aligned = whisperx.align(result["segments"], alignment_model, metadata, audio, device, return_char_alignments=True) print("Detecting speakers with WhisperX...") diarization_model = whisperx.DiarizationPipeline(use_auth_token=token, device=device) diarization_segments = diarization_model(audio) speaker_map = assign_speakers(result_aligned["segments"], diarization_segments) for segment in result_aligned["segments"]: segment["speaker"] = speaker_map.get(segment["start"], "Unknown") segment.pop("chars", None) session_id = generate_session_id() session_dir = os.path.join("session_data", session_id) os.makedirs(session_dir, exist_ok=True) data, sr = load_audio_for_split(input_audio_file) for segment in result_aligned["segments"]: start_time = segment["start"] end_time = segment["end"] speaker = segment["speaker"] start_sample = int(start_time * sr) end_sample = int(end_time * sr) segment_audio = data[start_sample:end_sample] segment_filename = f"{session_id}-{start_time:.2f}-{end_time:.2f}-{speaker}.wav" segment_filepath = os.path.join(session_dir, segment_filename) sf.write(segment_filepath, segment_audio, sr) print(f"Saved segment: {segment_filepath}") transcript_path = os.path.join(session_dir, f"{session_id}_transcription.txt") with open(transcript_path, "w", encoding="utf-8") as f: for segment in result_aligned["segments"]: f.write(f"[{segment['start']} - {segment['end']}] (Speaker {segment['speaker']}): {segment['text']}\n") del model torch.cuda.empty_cache() print("Loading CrisperWhisper model...") device_str = "cuda:0" if torch.cuda.is_available() else "cpu" torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 """ Use local Crisper Whisper Model local_model_dir = "./CrisperWhisper_local" cw_model = AutoModelForSpeechSeq2Seq.from_pretrained( local_model_dir, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True ) cw_model.to(device_str) processor = AutoProcessor.from_pretrained(local_model_dir) """ hf_model_id = "nyrahealth/CrisperWhisper" cw_model = AutoModelForSpeechSeq2Seq.from_pretrained( hf_model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, token=token ) cw_model.to(device_str) processor = AutoProcessor.from_pretrained(hf_model_id, token=token) asr_pipeline = pipeline( "automatic-speech-recognition", model=cw_model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, chunk_length_s=30, batch_size=4, return_timestamps='word', torch_dtype=torch_dtype, device=0 if torch.cuda.is_available() else -1, generate_kwargs={"language": "en"} ) segments_cw = [] skipped_segments = [] segment_files = [f for f in os.listdir(session_dir) if f.endswith('.wav')] for seg_file in sorted(segment_files): match = re.match(r'^(\d+)-(\d+\.\d+)-(\d+\.\d+)-(.+)\.wav$', seg_file) if not match: continue seg_session_id = match.group(1) start_time = float(match.group(2)) end_time = float(match.group(3)) speaker = match.group(4) seg_path = os.path.join(session_dir, seg_file) print(f"Processing segment with CrisperWhisper: {seg_path}") try: cw_output = asr_pipeline(seg_path) cw_result = adjust_pauses_for_hf_pipeline_output(cw_output) except Exception as e: print(f"[Warning] CrisperWhisper error, skiped this segment: {seg_path}\nError Message: {e}") skipped_segments.append(seg_path) continue text = cw_result.get('text', '').strip() if not text: print(f"********** No text returned, skiped this segment: {seg_path} **********") skipped_segments.append(seg_path) continue chunks = cw_result.get('chunks', []) words_info = [] for i, chunk in enumerate(chunks): word_text = chunk['text'].strip() if not word_text: continue chunk_start, chunk_end = chunk['timestamp'] if chunk_start is None: if i == 0: chunk_start = 0.0 else: chunk_start = words_info[-1]['end'] - start_time if chunk_end is None: if i < len(chunks) - 1: next_chunk_start, _ = chunks[i+1]['timestamp'] if next_chunk_start is None: next_chunk_start = chunk_start chunk_end = next_chunk_start else: chunk_end = end_time - start_time word_start = round(start_time + chunk_start, 3) word_end = round(start_time + chunk_end, 3) words_info.append({ "word": word_text, "start": word_start, "end": word_end }) segment_entry = { "start": round(start_time, 3), "end": round(end_time, 3), "speaker": speaker, "text": text, "words": words_info } print(f"Post-processing: splitting segment by sentences...") split_segments = split_segment_by_sentences(segment_entry) segments_cw.extend(split_segments) segments_cw = sorted(segments_cw, key=lambda x: x["start"]) cw_json_path = os.path.join(session_dir, f"{session_id}_transcriptionCW.json") with open(cw_json_path, "w", encoding="utf-8") as f: json.dump({"segments": segments_cw}, f, ensure_ascii=False, indent=4) print(f"CrisperWhisper transcription saved to: {cw_json_path}") if skipped_segments: skipped_file = os.path.join(session_dir, "skipped_segments.txt") with open(skipped_file, "w", encoding="utf-8") as f: for s in sorted(skipped_segments): f.write(s + "\n") print(f"Skipped segments recorded in: {skipped_file}") return session_id if __name__ == "__main__": session = process_audio_file("/home/easgrad/shuweiho/workspace/volen/SATE_docker_test/input/454.mp3", num_speakers=2, device="cuda") print("Processing complete. Session ID:", session)