create app.py

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+import torch
+import torchaudio
+import re  # Import regex library
+
+# Load model and processor
+processor = WhisperProcessor.from_pretrained("openai/whisper-large-v2")
+model = WhisperForConditionalGeneration.from_pretrained("aiola/whisper-ner-v1").to("cuda")
+
+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 transcribe_and_recognize_entities(audio_file, prompt):
+    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)
+
+    signal = signal.cpu()  # Ensure signal is on CPU for processing
+    input_features = processor(signal, sampling_rate=target_sample_rate, return_tensors="pt").input_features
+
+    # Split the prompt into individual NER types and process each one
+    ner_types = prompt.split(',')
+    processed_ner_types = [unify_ner_text(ner_type.strip()) for ner_type in ner_types]
+    prompt = ", ".join(processed_ner_types)
+
+    print(f"Prompt after unify_ner_text: {prompt}")
+    prompt_ids = processor.get_prompt_ids(prompt, return_tensors="pt")
+    prompt_ids = prompt_ids.to("cuda")
+
+    predicted_ids = model.generate(
+        input_features.to("cuda"),
+        max_new_tokens=256,
+        prompt_ids=prompt_ids,
+        language='en',  # Ensure transcription is translated to English
+        generation_config=model.generation_config,
+    )
+    # slice only the output without the prompt itself at the start.
+    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+
+    # Determine the length of the prompt in the transcription
+    prompt_length_in_transcription = len(prompt)
+
+    # Slice the transcription to remove the prompt itself from the output
+    transcription = transcription[prompt_length_in_transcription + 1:]
+
+    return transcription
+
+# Define Gradio interface
+iface = gr.Interface(
+    fn=transcribe_and_recognize_entities,
+    inputs=[
+        gr.Audio(label="Upload Audio", type="filepath"),
+        gr.Textbox(label="Entity Recognition Prompt"),
+    ],
+    outputs=gr.Textbox(label="Transcription and Entities"),
+    title="Whisper-NER Demo",
+    description="Upload an audio file and enter entities to identify. The model will transcribe the audio and recognize entities."
+)
+
+# iface.launch()
+iface.launch(share=True)