disable whishper modes

ColoredCow · Nov 12, 2024 · 932b65f · 932b65f
1 parent 0b09a53
commit 932b65f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/app.py b/app.py
@@ -7,7 +7,7 @@
 from dotenv import load_dotenv
 import os
 from huggingface_hub import login
-from transcription import load_asr_model, translate_with_base_whisper, translate_audio
+from transcription import load_asr_model, translate_audio
 
 # Load the environment variables from the .env file
 load_dotenv()

diff --git a/transcription.py b/transcription.py
@@ -4,8 +4,6 @@
 
 import whisper
 
-whisper_model = whisper.load_model("small")
-
 # Load the Whisper model and processor from Hugging Face
 def load_asr_model(modelName):
     processor = WhisperProcessor.from_pretrained(modelName)
@@ -37,20 +35,22 @@ def translate_audio(file_path, model, processor, language):
 
     # Generate transcription using the fine-tuned model
     with torch.no_grad():
-        generated_tokens = model.generate(input_features, forced_decoder_ids=processor.get_decoder_prompt_ids(language="mr", task="translate"))
+        generated_tokens = model.generate(input_features, forced_decoder_ids=processor.get_decoder_prompt_ids(language=language, task="translate"))
         transcription = processor.decode(generated_tokens[0], skip_special_tokens=True)
 
     return transcription
 
 def translate_with_base_whisper(file_path, model, processor, language):
     # Transcribe using Whisper
+    whisper_model = whisper.load_model("small")
     result = whisper_model.transcribe(file_path, task="translate", language = language)
     transcription = result['text']
     print(f"Transcription: {transcription}")
     return transcription
 
 def transcribe_with_base_whisper(file_path, model, processor, language):
     # Transcribe using Whisper
+    whisper_model = whisper.load_model("small")
     result = whisper_model.transcribe(file_path, task="transcribe", language = language)
     transcription = result['text']
     print(f"Transcription: {transcription}")