import gradio as gr from transformers import WhisperProcessor, WhisperForConditionalGeneration import torch import torchaudio import spaces import re # Initialize devices device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load model and processor processor = WhisperProcessor.from_pretrained("aiola/whisper-ner-tag-and-mask-v1") model = WhisperForConditionalGeneration.from_pretrained("aiola/whisper-ner-tag-and-mask-v1") model = model.to(device) examples = [ [ "audio/sports.wav", "football-club, football-player, referee", False ], [ "audio/entertainment.wav", "movie, date, actor, tv-show, musician", True ], [ "audio/personal_info.wav", "address, name, phone-number", True ], [ "audio/672-122797-0026.wav", "biological-classification, desire, demographic-group, object-category, relationship-role, reflexive-pronoun, furniture-type", False ], [ "audio/672-122797-0027.wav", "action, emotional-resilience, comparative-path-characteristic, social-role", True ], [ "audio/672-122797-0024.wav", "health-warning, importance-indicator, event, sentiment", False ], [ "audio/672-122797-0048.wav", "weapon, emotional-state, household-chore, atmosphere-quality", False ], ] def unify_ner_text(text, symbols_to_replace=("/", " ", ":", "_")): """Process and standardize entity text by replacing certain symbols and normalizing spaces.""" text = " ".join(text.split()) for symbol in symbols_to_replace: text = text.replace(symbol, "-") return text.lower() def extract_entities_and_clean_text_fixed(text, ner_mask=False): entity_pattern = r"<(.*?)>(.*?)<\1>>" if not ner_mask else r"<(.*?)>>" entities = [] clean_text = [] current_pos = 0 # Iterate through the matches for entity tags for match in re.finditer(entity_pattern, text): # Add text before the entity to the clean text clean_text.append(text[current_pos:match.start()]) entity_type = match.group(1) entity_text = "-" if ner_mask else match.group(2) start_pos = len("".join(clean_text)) # Start position in the clean text end_pos = start_pos + len(entity_text) # Append the entity text to the clean text clean_text.append(entity_text) # Add the entity details to the list entities.append({ "entity": entity_type, "text": entity_text, "start": start_pos, "end": end_pos }) # Update the current position to the end of the match current_pos = match.end() # Append the remaining part of the text after the last entity clean_text.append(text[current_pos:]) # Join all parts of the clean text clean_text_str = "".join(clean_text) return clean_text_str, entities @spaces.GPU # This decorator ensures your function can use GPU on Hugging Face Spaces def transcribe_and_recognize_entities(audio_file, prompt, ner_mask=False): target_sample_rate = 16000 signal, sampling_rate = torchaudio.load(audio_file) resampler = torchaudio.transforms.Resample(orig_freq=sampling_rate, new_freq=target_sample_rate) signal = resampler(signal) if signal.ndim == 2: signal = torch.mean(signal, dim=0) input_features = processor(signal, sampling_rate=target_sample_rate, return_tensors="pt").input_features input_features = input_features.to(device) ner_types = prompt.split(',') processed_ner_types = [unify_ner_text(ner_type.strip()) for ner_type in ner_types] prompt = ", ".join(processed_ner_types) if ner_mask: prompt = f"<|mask|>{prompt}" print(f"Prompt after unify_ner_text: {prompt}") prompt_ids = processor.get_prompt_ids(prompt, return_tensors="pt") prompt_ids = prompt_ids.to(device) predicted_ids = model.generate( input_features, max_new_tokens=256, prompt_ids=prompt_ids, language='en', generation_config=model.generation_config, ) transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0] clean_text_fixed, extracted_entities_fixed = extract_entities_and_clean_text_fixed(transcription, ner_mask=ner_mask) return transcription, {"text": clean_text_fixed, "entities": extracted_entities_fixed} with gr.Blocks(title="WhisperNER v1") as demo: gr.Markdown( """ # 🔥 Whisper-NER: ASR with zero-shot NER WhisperNER is a unified model for automatic speech recognition (ASR) and named entity recognition (NER), with zero-shot capabilities. The WhisperNER model is designed as a strong base model for the downstream task of ASR with NER, and can be fine-tuned on specific datasets for improved performance. The [aiola/whisper-ner-tag-and-mask-v1](https://huggingface.co/aiola/whisper-ner-tag-and-mask-v1) model was finetuned from the [aiola/whisper-ner-v1](https://huggingface.co/aiola/whisper-ner-v1) checkpoint using the NuNER dataset to perform joint audio transcription and NER tagging or NER masking. The model was not trained on PII specific datasets, hence can perform general and open type entity masking. It should be further finetuned in order to be used for PII detection. The model was trained and evaluated only on English data. Check out the paper for full details. ## Links * 📄 Paper: [WhisperNER: Unified Open Named Entity and Speech Recognition](https://arxiv.org/abs/2409.08107) * 🤗 [WhisperNER model collection](https://huggingface.co/collections/aiola/whisperner-6723f14506f3662cf3a73df2) * 💻 Code: https://github.com/aiola-lab/whisper-ner """ ) with gr.Row() as row1: with gr.Column() as col1: audio_input = gr.Audio(value=examples[0][0], label="Audio Example", type="filepath") with gr.Column() as col2: label_input = gr.Textbox(label="Entity Labels", value=examples[0][1]) ner_mask = gr.Checkbox( value=examples[0][2], label="Entity Mask", info="Mask or tag entities in the transcription.", scale=0, ) submit_btn = gr.Button("Submit") gr.Markdown("## Output") with gr.Row() as row3: transcript_output = gr.Textbox(label="Transcription and Entities") with gr.Row() as row4: highlighted_text_output = gr.HighlightedText(label="Predicted Highlighted Entities") examples = gr.Examples( examples, fn=transcribe_and_recognize_entities, inputs=[audio_input, label_input, ner_mask], outputs=[transcript_output, highlighted_text_output], cache_examples=True, run_on_click=True, ) # Submitting label_input.submit( fn=transcribe_and_recognize_entities, inputs=[audio_input, label_input, ner_mask], outputs=[transcript_output, highlighted_text_output], ) submit_btn.click( fn=transcribe_and_recognize_entities, inputs=[audio_input, label_input, ner_mask], outputs=[transcript_output, highlighted_text_output], ) demo.launch()