Add support for multiple languages (#18)

* Add support for multi languages.

TranscriberModels.py

@@ -16,29 +16,47 @@ def get_model(use_api: bool, model: str = None):
 class WhisperTranscriber:
  def __init__(self, model: str = 'tiny'):
  model_filename = model + '.en.pt'
+ self.lang = 'en'
+ self.model = model
 
  if not os.path.isfile(model_filename):
  print(f'Could not find the model file: {model_filename}')
  print(f'Download the model file and add it to the directory: \
  {os.getcwd()}')
- print('small model is available at: \
+ print('tiny multi-lingual model is available at: \
+ https://drive.google.com/file/d/1M4AFutTmQROaE9xk2jPc5Y4oFRibHhEh/view?usp=drive_link')
+ print('small english model is available at: \
  https://drive.google.com/file/d/1vhtoZCwfYGi5C4jK1r-QVr5GobSBnKiH/view?usp=drive_link')
- print('base model is available at: \
+ print('small multi-lingual model is available at: \
+ https://drive.google.com/file/d/1bl8er_st8WPZKPWVeYMNlaUi9IzR3jEZ/view?usp=drive_link')
+ print('base english model is available at: \
  https://drive.google.com/file/d/1E44DVjpfZX8tSrSagaDJXU91caZOkwa6/view?usp=drive_link')
+ print('base multi-lingual model is available at: \
+ https://drive.google.com/file/d/1UcqU_D0cPFqq_nckSfstMBfogFsvR-KR/view?usp=drive_link')
  exit()
-
- self.audio_model = whisper.load_model(os.path.join(os.getcwd(),
- model_filename))
+ self.model_filename = os.path.join(os.getcwd(), model_filename)
+ self.audio_model = whisper.load_model(self.model_filename)
  print(f'[INFO] Whisper using GPU: {str(torch.cuda.is_available())}')
 
  def get_transcription(self, wav_file_path):
  try:
- result = self.audio_model.transcribe(wav_file_path, fp16=torch.cuda.is_available())
+ result = self.audio_model.transcribe(wav_file_path,
+ fp16=torch.cuda.is_available(), language=self.lang)
  except Exception as exception:
  print(exception)
  return ''
  return result['text'].strip()
 
+ def change_lang(self, lang: str):
+ self.lang = lang
+ self.load_model()
+
+ def load_model(self):
+ if self.lang == "en":
+ self.audio_model = whisper.load_model(os.path.join(os.getcwd(), self.model + 'en.pt'))
+ else:
+ self.audio_model = whisper.load_model(os.path.join(os.getcwd(), self.model + '.pt'))
+
 
 class APIWhisperTranscriber:
  def __init__(self):

language.py

@@ -0,0 +1,101 @@
+LANGUAGES = {
+ "en": "english",
+ "zh": "chinese",
+ "de": "german",
+ "es": "spanish",
+ "ru": "russian",
+ "ko": "korean",
+ "fr": "french",
+ "ja": "japanese",
+ "pt": "portuguese",
+ "tr": "turkish",
+ "pl": "polish",
+ "ca": "catalan",
+ "nl": "dutch",
+ "ar": "arabic",
+ "sv": "swedish",
+ "it": "italian",
+ "id": "indonesian",
+ "hi": "hindi",
+ "fi": "finnish",
+ "vi": "vietnamese",
+ "iw": "hebrew",
+ "uk": "ukrainian",
+ "el": "greek",
+ "ms": "malay",
+ "cs": "czech",
+ "ro": "romanian",
+ "da": "danish",
+ "hu": "hungarian",
+ "ta": "tamil",
+ "no": "norwegian",
+ "th": "thai",
+ "ur": "urdu",
+ "hr": "croatian",
+ "bg": "bulgarian",
+ "lt": "lithuanian",
+ "la": "latin",
+ "mi": "maori",
+ "ml": "malayalam",
+ "cy": "welsh",
+ "sk": "slovak",
+ "te": "telugu",
+ "fa": "persian",
+ "lv": "latvian",
+ "bn": "bengali",
+ "sr": "serbian",
+ "az": "azerbaijani",
+ "sl": "slovenian",
+ "kn": "kannada",
+ "et": "estonian",
+ "mk": "macedonian",
+ "br": "breton",
+ "eu": "basque",
+ "is": "icelandic",
+ "hy": "armenian",
+ "ne": "nepali",
+ "mn": "mongolian",
+ "bs": "bosnian",
+ "kk": "kazakh",
+ "sq": "albanian",
+ "sw": "swahili",
+ "gl": "galician",
+ "mr": "marathi",
+ "pa": "punjabi",
+ "si": "sinhala",
+ "km": "khmer",
+ "sn": "shona",
+ "yo": "yoruba",
+ "so": "somali",
+ "af": "afrikaans",
+ "oc": "occitan",
+ "ka": "georgian",
+ "be": "belarusian",
+ "tg": "tajik",
+ "sd": "sindhi",
+ "gu": "gujarati",
+ "am": "amharic",
+ "yi": "yiddish",
+ "lo": "lao",
+ "uz": "uzbek",
+ "fo": "faroese",
+ "ht": "haitian creole",
+ "ps": "pashto",
+ "tk": "turkmen",
+ "nn": "nynorsk",
+ "mt": "maltese",
+ "sa": "sanskrit",
+ "lb": "luxembourgish",
+ "my": "myanmar",
+ "bo": "tibetan",
+ "tl": "tagalog",
+ "mg": "malagasy",
+ "as": "assamese",
+ "tt": "tatar",
+ "haw": "hawaiian",
+ "ln": "lingala",
+ "ha": "hausa",
+ "ba": "bashkir",
+ "jw": "javanese",
+ "su": "sundanese",
+}

main.py

@@ -15,6 +15,7 @@
 import interactions
 import ui
 from requests.exceptions import ConnectionError
+from language import LANGUAGES
 
 
 def main():
@@ -27,11 +28,17 @@ def main():
  cmd_args.add_argument('-m', '--model', action='store', choices=['tiny', 'base', 'small'],
  default='tiny',
  help='Specify the model to use for transcription.'
- '\nBy default tiny model is part of the install.'
- '\nbase model has to be downloaded from the link \
+ '\nBy default tiny english model is part of the install.'
+ '\ntiny multi-lingual model has to be downloaded from the link \
+ https://drive.google.com/file/d/1M4AFutTmQROaE9xk2jPc5Y4oFRibHhEh/view?usp=drive_link'
+ '\nbase english model has to be downloaded from the link \
  https://drive.google.com/file/d/1E44DVjpfZX8tSrSagaDJXU91caZOkwa6/view?usp=drive_link'
- '\nsmall model has to be downloaded from the link \
+ '\nbase multi-lingual model has to be downloaded from the link \
+ https://drive.google.com/file/d/1UcqU_D0cPFqq_nckSfstMBfogFsvR-KR/view?usp=drive_link'
+ '\nsmall english model has to be downloaded from the link \
  https://drive.google.com/file/d/1vhtoZCwfYGi5C4jK1r-QVr5GobSBnKiH/view?usp=drive_link'
+ '\nsmall multi-lingual model has to be downloaded from the link \
+ https://drive.google.com/file/d/1bl8er_st8WPZKPWVeYMNlaUi9IzR3jEZ/view?usp=drive_link'
  '\nOpenAI has more models besides the ones specified above.'
  '\nThose models are prohibitive to use on local machines because \
  of memory requirements.')
@@ -65,7 +72,8 @@ def main():
  freeze_button = ui_components[4]
  copy_button = ui_components[5]
  save_file_button = ui_components[6]
- transcript_button = ui_components[7]
+ lang_combobox = ui_components[7]
+ transcript_button = ui_components[8]
 
  audio_queue = queue.Queue()
 
@@ -136,6 +144,8 @@ def set_transcript_state():
  {update_interval_slider.get()} \
  seconds")
 
+ lang_combobox.configure(command=model.change_lang)
+
  ui.update_transcript_ui(transcriber, transcript_textbox)
  ui.update_response_ui(responder, response_textbox, update_interval_slider_label,
  update_interval_slider, freeze_state)

ui.py

@@ -1,6 +1,7 @@
 import customtkinter as ctk
 import AudioTranscriber
 import prompts
+from language import LANGUAGES
 
 
 def write_in_textbox(textbox: ctk.CTkTextbox, text: str):
@@ -91,11 +92,14 @@ def create_ui_components(root):
  save_file_button = ctk.CTkButton(root, text="Save Audio Transcript to File", command=None)
  save_file_button.grid(row=3, column=0, padx=10, pady=3, sticky="nsew")
 
+ lang_combobox = ctk.CTkOptionMenu(root, values=list(LANGUAGES.values()))
+ lang_combobox.grid(row=4, column=1, padx=200, pady=10, sticky="nsew")
+
  transcript_button = ctk.CTkButton(root, text="Pause Transcript", command=None)
  transcript_button.grid(row=4, column=0, padx=10, pady=3, sticky="nsew")
 
  # Order of returned components is important.
  # Add new components to the end
  return [transcript_textbox, response_textbox, update_interval_slider,
  update_interval_slider_label, freeze_button, copy_button,
- save_file_button, transcript_button]
+ save_file_button, lang_combobox, transcript_button]