import os
import gc
import sys
import time
import torch
import spaces
import torchaudio
import numpy as np
from df.enhance import enhance, init_df
from dotenv import load_dotenv
load_dotenv()
from scipy.signal import resample
from pyannote.audio import Pipeline
from difflib import SequenceMatcher
from transformers import Wav2Vec2ForSequenceClassification, AutoFeatureExtractor, Wav2Vec2ForCTC, AutoProcessor, AutoTokenizer, AutoModelForSeq2SeqLM
import logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)
class ChunkedTranscriber:
def __init__(self, chunk_size=30, overlap=5, sample_rate=16000):
self.chunk_size = chunk_size
self.overlap = overlap
self.sample_rate = sample_rate
self.previous_text = ""
self.previous_lang = None
self.speaker_diarization_pipeline = self.load_speaker_diarization_pipeline()
def load_speaker_diarization_pipeline(self):
"""
Load the pre-trained speaker diarization pipeline from pyannote-audio.
"""
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=os.getenv("HF_TOKEN"))
return pipeline
@spaces.GPU(duration=180)
def diarize_audio(self, audio_path):
"""
Perform speaker diarization on the input audio.
"""
diarization_result = self.speaker_diarization_pipeline({"uri": "audio", "audio": audio_path})
return diarization_result
def load_lid_mms(self):
model_id = "facebook/mms-lid-256"
processor = AutoFeatureExtractor.from_pretrained(model_id)
model = Wav2Vec2ForSequenceClassification.from_pretrained(model_id)
return processor, model
@spaces.GPU(duration=180)
def language_identification(self, model, processor, chunk, device="cuda"):
inputs = processor(chunk, sampling_rate=16_000, return_tensors="pt")
model.to(device)
inputs.to(device)
with torch.no_grad():
outputs = model(**inputs).logits
lang_id = torch.argmax(outputs, dim=-1)[0].item()
detected_lang = model.config.id2label[lang_id]
del model
del inputs
torch.cuda.empty_cache()
gc.collect()
return detected_lang
def load_mms(self) :
model_id = "facebook/mms-1b-all"
processor = AutoProcessor.from_pretrained(model_id)
model = Wav2Vec2ForCTC.from_pretrained(model_id)
return model, processor
@spaces.GPU(duration=180)
def mms_transcription(self, model, processor, chunk, device="cuda"):
inputs = processor(chunk, sampling_rate=16_000, return_tensors="pt")
model.to(device)
inputs.to(device)
with torch.no_grad():
outputs = model(**inputs).logits
ids = torch.argmax(outputs, dim=-1)[0]
transcription = processor.decode(ids)
del model
del inputs
torch.cuda.empty_cache()
gc.collect()
return transcription
def load_T2T_translation_model(self) :
model_id = "facebook/nllb-200-distilled-600M"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
return model, tokenizer
@spaces.GPU(duration=180)
def text2text_translation(self, translation_model, translation_tokenizer, transcript, device="cuda"):
# model, tokenizer = load_translation_model()
tokenized_inputs = translation_tokenizer(transcript, return_tensors='pt')
translation_model.to(device)
tokenized_inputs.to(device)
translated_tokens = translation_model.generate(**tokenized_inputs,
forced_bos_token_id=translation_tokenizer.convert_tokens_to_ids("eng_Latn"),
max_length=100)
del translation_model
del tokenized_inputs
torch.cuda.empty_cache()
gc.collect()
return translation_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
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
def merge_close_segments(self, results):
"""
Merge segments that are close in time and have the same language
"""
if not results:
return results
merged = []
current = results[0]
for next_segment in results[1:]:
# Skip empty segments
if not next_segment['text'].strip():
continue
# If segments are in the same language and close in time
if (current['detected_language'] == next_segment['detected_language'] and
abs(next_segment['start_time'] - current['end_time']) <= self.overlap):
# Merge the segments
current['text'] = current['text'] + ' ' + next_segment['text']
current['end_time'] = next_segment['end_time']
if 'translated' in current and 'translated' in next_segment:
current['translated'] = current['translated'] + ' ' + next_segment['translated']
else:
if current['text'].strip(): # Only add non-empty segments
merged.append(current)
current = next_segment
if current['text'].strip(): # Add the last segment if non-empty
merged.append(current)
return merged
def clean_overlapping_text(self, current_text, prev_text, current_lang, prev_lang, min_overlap=3):
"""
Improved text cleaning with language awareness and better sentence boundary handling
"""
if not prev_text or not current_text:
return current_text
# If languages are different, don't try to merge
if prev_lang and current_lang and prev_lang != current_lang:
return current_text
# Split into words
prev_words = prev_text.split()
curr_words = current_text.split()
if len(prev_words) < 2 or len(curr_words) < 2:
return current_text
# Find matching sequences at the end of prev_text and start of current_text
matcher = SequenceMatcher(None, prev_words, curr_words)
matches = list(matcher.get_matching_blocks())
# Look for significant overlaps
best_overlap = 0
overlap_size = 0
for match in matches:
# Check if the match is at the start of current text
if match.b == 0 and match.size >= min_overlap:
if match.size > overlap_size:
best_overlap = match.size
overlap_size = match.size
if best_overlap > 0:
# Remove overlapping content while preserving sentence integrity
cleaned_words = curr_words[best_overlap:]
if not cleaned_words: # If everything was overlapping
return ""
return ' '.join(cleaned_words).strip()
return current_text
def process_chunk(self, chunk_data, mms_model, mms_processor, translation_model=None, translation_tokenizer=None):
"""
Process chunk with improved language handling
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
try:
# Language detection
lid_processor, lid_model = self.load_lid_mms()
lid_lang = self.language_identification(lid_model, lid_processor, chunk_data['chunk'])
# Configure processor
mms_processor.tokenizer.set_target_lang(lid_lang)
mms_model.load_adapter(lid_lang)
# Transcribe
inputs = mms_processor(chunk_data['chunk'], sampling_rate=self.sample_rate, return_tensors="pt")
inputs = inputs.to(device)
mms_model = mms_model.to(device)
with torch.no_grad():
outputs = mms_model(**inputs).logits
ids = torch.argmax(outputs, dim=-1)[0]
transcription = mms_processor.decode(ids)
# Clean overlapping text with language awareness
cleaned_transcription = self.clean_overlapping_text(
transcription,
self.previous_text,
lid_lang,
self.previous_lang,
min_overlap=3
)
# Update previous state
self.previous_text = transcription
self.previous_lang = lid_lang
if not cleaned_transcription.strip():
return None
result = {
'start_time': chunk_data['start_time'],
'end_time': chunk_data['end_time'],
'text': cleaned_transcription,
'detected_language': lid_lang
}
# Handle translation
if translation_model and translation_tokenizer and cleaned_transcription.strip():
translation = self.text2text_translation(
translation_model,
translation_tokenizer,
cleaned_transcription
)
result['translated'] = translation
return result
except Exception as e:
print(f"Error processing chunk: {str(e)}")
return None
finally:
torch.cuda.empty_cache()
gc.collect()
def translate_text(self, text, translation_model, translation_tokenizer, device):
"""
Translate cleaned text using the provided translation model.
"""
tokenized_inputs = translation_tokenizer(text, return_tensors='pt')
tokenized_inputs = tokenized_inputs.to(device)
translation_model = translation_model.to(device)
translated_tokens = translation_model.generate(
**tokenized_inputs,
forced_bos_token_id=translation_tokenizer.convert_tokens_to_ids("eng_Latn"),
max_length=100
)
translation = translation_tokenizer.batch_decode(
translated_tokens,
skip_special_tokens=True
)[0]
del translation_model
del tokenized_inputs
torch.cuda.empty_cache()
gc.collect()
return translation
def audio_denoising():
model, df_state = init_df()
enhanced_audio = enhance(model, df_state, noisy_audio)
return enhanced_audio
def transcribe_audio(self, audio_path, translate=False):
"""
Main transcription function with improved segment merging
"""
# Perform speaker diarization
diarization_result = self.diarize_audio(audio_path)
# Extract speaker segments
speaker_segments = []
for turn, _, speaker in diarization_result.itertracks(yield_label=True):
speaker_segments.append({
'start_time': turn.start,
'end_time': turn.end,
'speaker': speaker
})
audio = self.load_audio(audio_path)
chunks = self.preprocess_audio(audio)
mms_model, mms_processor = self.load_mms()
translation_model, translation_tokenizer = None, None
if translate:
translation_model, translation_tokenizer = self.load_T2T_translation_model()
# Process chunks
results = []
for chunk_data in chunks:
result = self.process_chunk(
chunk_data,
mms_model,
mms_processor,
translation_model,
translation_tokenizer
)
if result:
for segment in speaker_segments:
if int(segment['start_time']) <= int(chunk_data['start_time']) < int(segment['end_time']):
result['speaker'] = segment['speaker']
break
results.append(result)
# results.append(result)
# Merge close segments and clean up
merged_results = self.merge_close_segments(results)
_translation = ""
_output = ""
for res in merged_results:
_translation+=res['translated']
_output+=f"{res['start_time']}-{res['end_time']} - Speaker: {res['speaker'].split('_')[1]} - Language: {res['detected_language']}\n Text: {res['text']}\n Translation: {res['translated']}\n\n"
logger.info(f"\n\n TRANSLATION: {_translation}")
return _translation, _output
def load_audio(self, audio_path):
"""
Load and preprocess audio file.
"""
waveform, sample_rate = torchaudio.load(audio_path)
# Convert to mono if stereo
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)
return waveform.float()