update the app.py

app.py

@@ -2,11 +2,14 @@ import gradio as gr
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 import torch
 import torchaudio
+import spaces
+
+# Initialize devices
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load model and processor
 processor = WhisperProcessor.from_pretrained("openai/whisper-large-v2")
 model = WhisperForConditionalGeneration.from_pretrained("aiola/whisper-ner-v1")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
 def unify_ner_text(text, symbols_to_replace=("/", " ", ":", "_")):
@@ -16,7 +19,7 @@ def unify_ner_text(text, symbols_to_replace=("/", " ", ":", "_")):
         text = text.replace(symbol, "-") 
     return text.lower()
 
-
+@spaces.GPU  # This decorator ensures your function can use GPU on Hugging Face Spaces
 def transcribe_and_recognize_entities(audio_file, prompt):
     target_sample_rate = 16000
     signal, sampling_rate = torchaudio.load(audio_file)
@@ -25,12 +28,9 @@ def transcribe_and_recognize_entities(audio_file, prompt):
     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
     input_features = input_features.to(device)
 
-
-    # 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)
@@ -43,31 +43,21 @@ def transcribe_and_recognize_entities(audio_file, prompt):
         input_features,
         max_new_tokens=256,
         prompt_ids=prompt_ids,
-        language='en',  # Ensure transcription is translated to English
+        language='en',
         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:]
+    transcription = transcription[prompt_length_in_transcription + 1:]  # Remove the prompt
 
     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"),
-    ],
+    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)