Merge branch 'main' of https://huggingface.co/spaces/SoybeanMilk/whisper-webui-translate

1.Thank you SoybeanMilk for assisting in the development and integration of the excellent ALMA translation model.

2.Add the missing packages for the GPTQ version of ALMA in the requirements.txt.

3.Include the 7B version of the ALMA model in addition to the 13B version in the web UI.

4.Write the ReadMe document for the Translate Model, containing a brief introduction and explanation of each translation model used in this project.

app.py

@@ -40,7 +40,7 @@ from src.whisper.whisperFactory import create_whisper_container
 from src.translation.translationModel import TranslationModel
 from src.translation.translationLangs import (TranslationLang,
                                               _TO_LANG_CODE_WHISPER, get_lang_whisper_names, get_lang_from_whisper_name, get_lang_from_whisper_code, 
-                                              get_lang_nllb_names, get_lang_from_nllb_name, get_lang_m2m100_names, get_lang_from_m2m100_name)
+                                              get_lang_nllb_names, get_lang_from_nllb_name, get_lang_m2m100_names, get_lang_from_m2m100_name, sort_lang_by_whisper_codes)
 import shutil
 import zhconv
 import tqdm
@@ -773,8 +773,15 @@ class WhisperTranscriber:
             self.diarization = None
 
 def create_ui(app_config: ApplicationConfig):
+    translateModelMd: str = None
     optionsMd: str = None
     readmeMd: str = None
+    try:
+        translateModelPath = pathlib.Path("docs/translateModel.md")
+        with open(translateModelPath, "r", encoding="utf-8") as translateModelFile:
+            translateModelMd = translateModelFile.read()
+    except Exception as e:
+        print("Error occurred during read translateModel.md file: ", str(e))
     try:
         optionsPath = pathlib.Path("docs/options.md")
         with open(optionsPath, "r", encoding="utf-8") as optionsFile:
@@ -819,16 +826,6 @@ def create_ui(app_config: ApplicationConfig):
         uiDescription += "\n\n" + "Max audio file length: " + str(app_config.input_audio_max_duration) + " s"
 
     uiArticle = "Read the [documentation here](https://gitlab.com/aadnk/whisper-webui/-/blob/main/docs/options.md)."
-    uiArticle += "\n\nWhisper's Task 'translate' only implements the functionality of translating other languages into English. "
-    uiArticle += "OpenAI does not guarantee translations between arbitrary languages. In such cases, you can choose to use the Translation Model to implement the translation task. "
-    uiArticle += "However, it's important to note that the Translation Model runs slowly(in CPU), and the completion time may be twice as long as usual. "
-    uiArticle += "\n\nThe larger the parameters of the Translation model, the better its performance is expected to be. "
-    uiArticle += "However, it also requires higher computational resources, making it slower to operate. "
-    uiArticle += "On the other hand, the version converted from ct2 ([CTranslate2](https://opennmt.net/CTranslate2/guides/transformers.html)) requires lower resources and operates at a faster speed."
-    uiArticle += "\n\nCurrently, enabling `Highlight Words` timestamps cannot be used in conjunction with Translation Model translation "
-    uiArticle += "because Highlight Words will split the source text, and after translation, it becomes a non-word-level string. "
-    uiArticle += "\n\nThe 'mt5-zh-ja-en-trimmed' model is finetuned from Google's 'mt5-base' model. "
-    uiArticle += "This model has a relatively good translation speed, but it only supports three languages: Chinese, Japanese, and English. "
 
     whisper_models = app_config.get_model_names("whisper")
     nllb_models = app_config.get_model_names("nllb")
@@ -854,7 +851,7 @@ def create_ui(app_config: ApplicationConfig):
     }
     common_ALMA_inputs = lambda : {
         gr.Dropdown(label="ALMA - Model (for translate)", choices=ALMA_models, elem_id="ALMAModelName"),
-        gr.Dropdown(label="ALMA - Language", choices=sorted(get_lang_m2m100_names(["en", "ja", "zh"])), elem_id="ALMALangName"),
+        gr.Dropdown(label="ALMA - Language", choices=sort_lang_by_whisper_codes(["en", "de", "cs", "is", "ru", "zh", "ja"]), elem_id="ALMALangName"),
     }
     
     common_translation_inputs = lambda : {
@@ -944,8 +941,10 @@ def create_ui(app_config: ApplicationConfig):
                         simpleInputDict.update(common_translation_inputs())
             with gr.Column():
                 simpleOutput = common_output()
-        with gr.Accordion("Article"):
-            gr.Markdown(uiArticle)
+        gr.Markdown(uiArticle)
+        if translateModelMd is not None:
+            with gr.Accordion("docs/translateModel.md", open=False):    
+                gr.Markdown(translateModelMd)
         if optionsMd is not None:
             with gr.Accordion("docs/options.md", open=False):    
                 gr.Markdown(optionsMd)
@@ -1056,7 +1055,7 @@ def create_ui(app_config: ApplicationConfig):
         print("Queue mode enabled (concurrency count: " + str(app_config.queue_concurrency_count) + ")")
     else:
         print("Queue mode disabled - progress bars will not be shown.")
-   
+
     demo.launch(inbrowser=app_config.autolaunch, share=app_config.share, server_name=app_config.server_name, server_port=app_config.server_port)
     
     # Clean up
@@ -1138,6 +1137,16 @@ if __name__ == '__main__':
     # updated_config.autolaunch = True
     # updated_config.auto_parallel = False
     # updated_config.save_downloaded_files = True
+    
+    try:
+        if torch.cuda.is_available():
+            deviceId = torch.cuda.current_device()
+            totalVram = torch.cuda.get_device_properties(deviceId).total_memory
+            if totalVram/(1024*1024*1024) <= 4: #VRAM <= 4 GB
+                updated_config.vad_process_timeout = 0
+    except Exception as e:
+        print(traceback.format_exc())
+        print("Error detect vram: " + str(e))
 
     if (threads := args.pop("threads")) > 0:
         torch.set_num_threads(threads)

config.json5

@@ -193,10 +193,15 @@
       }
     ],
     "ALMA": [
+      {
+        "name": "ALMA-7B-GPTQ/TheBloke",
+        "url": "TheBloke/ALMA-7B-GPTQ",
+        "type": "huggingface"
+      },
       {
         "name": "ALMA-13B-GPTQ/TheBloke",
         "url": "TheBloke/ALMA-13B-GPTQ",
-        "type": "huggingface", 
+        "type": "huggingface"
       },
     ]
   },

docs/options.md

@@ -1,4 +1,4 @@
-# Standard Options
+# Standard Options
 To transcribe or translate an audio file, you can either copy an URL from a website (all [websites](https://github.com/yt-dlp/yt-dlp/blob/master/supportedsites.md) 
 supported by YT-DLP will work, including YouTube). Otherwise, upload an audio file (choose "All Files (*.*)" 
 in the file selector to select any file type, including video files) or use the microphone.
@@ -154,29 +154,29 @@ The minimum number of speakers for Pyannote to detect.
 The maximum number of speakers for Pyannote to detect.
 
 ## Repetition Penalty
-- ctranslate2: repetition_penalty
+- ctranslate2: repetition_penalty  
 This parameter only takes effect in [faster-whisper (ctranslate2)](https://github.com/SYSTRAN/faster-whisper/issues/478).
 Penalty applied to the score of previously generated tokens (set > 1 to penalize).
 
 ## No Repeat Ngram Size
-- ctranslate2: no_repeat_ngram_size
+- ctranslate2: no_repeat_ngram_size  
 This parameter only takes effect in [faster-whisper (ctranslate2)](https://github.com/SYSTRAN/faster-whisper/issues/478).
 Prevent repetitions of ngrams with this size (set 0 to disable).
 
 ## Translation - Batch Size
-- transformers: batch_size
+- transformers: batch_size  
 When the pipeline will use DataLoader (when passing a dataset, on GPU for a Pytorch model), the size of the batch to use, for inference this is not always beneficial.
-- ctranslate2: max_batch_size
+- ctranslate2: max_batch_size  
 The maximum batch size.
 
 ## Translation - No Repeat Ngram Size
-- transformers: no_repeat_ngram_size
+- transformers: no_repeat_ngram_size  
 Value that will be used by default in the generate method of the model for no_repeat_ngram_size. If set to int > 0, all ngrams of that size can only occur once.
-- ctranslate2: no_repeat_ngram_size
+- ctranslate2: no_repeat_ngram_size  
 Prevent repetitions of ngrams with this size (set 0 to disable).
 
 ## Translation - Num Beams
-- transformers: num_beams
+- transformers: num_beams  
 Number of beams for beam search that will be used by default in the generate method of the model. 1 means no beam search.
-- ctranslate2: beam_size
+- ctranslate2: beam_size  
 Beam size (1 for greedy search).

docs/translateModel.md

@@ -0,0 +1,112 @@
+
+# Describe
+
+The `translate` task in `Whisper` only supports translating other languages `into English`. `OpenAI` does not guarantee translations between arbitrary languages. In such cases, you can opt to use the Translation Model for translation tasks. However, it's important to note that the `Translation Model runs very slowly on CPU`, and the completion time may be twice as long as usual. It is recommended to run the Translation Model on devices with `GPUs` for better performance.  
+
+The larger the parameters of the Translation model, the better its translation capability is expected. However, this also requires higher computational resources and slower running speed.  
+
+Currently, when the `Highlight Words timestamps` option is enabled in the Whisper `Word Timestamps options`, it cannot be used simultaneously with the Translation Model. This is because Highlight Words splits the source text, and after translation, it becomes a non-word-level string.  
+
+
+# Translation Model
+
+The required VRAM is provided for reference and may not apply to everyone. If the model's VRAM requirement exceeds the available capacity of the system, the model will operate on the CPU, resulting in significantly longer execution times.  
+
+[CTranslate2](https://opennmt.net/CTranslate2/guides/transformers.html) is a C++ and Python library for efficient inference with Transformer models. Models converted from CTranslate2 can run with lower resources and faster speed. Encoder-decoder models currently supported: Transformer base/big, M2M-100, NLLB, BART, mBART, Pegasus, T5, Whisper.  
+
+## M2M100
+
+2M100 is a multilingual translation model introduced by Facebook AI in October 2020. It supports arbitrary translation among 101 languages. The paper is titled "`Beyond English-Centric Multilingual Machine Translation`" ([arXiv:2010.11125](https://arxiv.org/abs/2010.11125)).  
+
+| Name | Parameters | Size | type/quantize | Required VRAM |
+|------|------------|------|---------------|---------------|
+| [facebook/m2m100_418M](https://huggingface.co/facebook/m2m100_418M) | 480M | 1.94 GB | float32 | ≈2 GB |
+| [facebook/m2m100_1.2B](https://huggingface.co/facebook/m2m100_1.2B) | 1.2B | 4.96 GB | float32 | ≈5 GB |
+| [facebook/m2m100-12B-last-ckpt](https://huggingface.co/facebook/m2m100-12B-last-ckpt) | 12B | 47.2 GB | float32 | N/A |
+
+## M2M100-CTranslate2
+
+| Name | Parameters | Size | type/quantize | Required VRAM |
+|------|------------|------|---------------|---------------|
+| [michaelfeil/ct2fast-m2m100_418M](https://huggingface.co/michaelfeil/ct2fast-m2m100_418M) | 480M | 970 MB | float16 | ≈0.6 GB |
+| [michaelfeil/ct2fast-m2m100_1.2B](https://huggingface.co/michaelfeil/ct2fast-m2m100_1.2B) | 1.2B | 2.48 GB | float16 | ≈1.3 GB |
+| [michaelfeil/ct2fast-m2m100-12B-last-ckpt](https://huggingface.co/michaelfeil/ct2fast-m2m100-12B-last-ckpt) | 12B | 23.6 GB | float16 | N/A |
+
+## NLLB-200
+
+NLLB-200 is a multilingual translation model introduced by Meta AI in July 2022. It supports arbitrary translation among 202 languages. The paper is titled "`No Language Left Behind: Scaling Human-Centered Machine Translation`" ([arXiv:2207.04672](https://arxiv.org/abs/2207.04672)).  
+
+| Name | Parameters | Size | type/quantize | Required VRAM |
+|------|------------|------|---------------|---------------|
+| [facebook/nllb-200-distilled-600M](https://huggingface.co/facebook/nllb-200-distilled-600M) | 600M | 2.46 GB | float32 | ≈2.5 GB |
+| [facebook/nllb-200-distilled-1.3B](https://huggingface.co/facebook/nllb-200-distilled-1.3B) | 1.3B | 5.48 GB | float32 | ≈5.9 GB |
+| [facebook/nllb-200-1.3B](https://huggingface.co/facebook/nllb-200-1.3B) | 1.3B | 5.48 GB | float32 | 5.8 GB |
+| [facebook/nllb-200-3.3B](https://huggingface.co/facebook/nllb-200-3.3B) | 3.3B | 17.58 GB | float32 | 13.4 GB |
+
+## NLLB-200-CTranslate2
+
+| Name | Parameters | Size | type/quantize | Required VRAM |
+|------|------------|------|---------------|---------------|
+| [michaelfeil/ct2fast-nllb-200-distilled-1.3B](https://huggingface.co/michaelfeil/ct2fast-nllb-200-distilled-1.3B) | 1.3B | 1.38 GB | int8_float16 | ≈1.3 GB |
+| [michaelfeil/ct2fast-nllb-200-3.3B](https://huggingface.co/michaelfeil/ct2fast-nllb-200-3.3B) | 3.3B | 3.36 GB | int8_float16 | ≈3.2 GB |
+| [JustFrederik/nllb-200-1.3B-ct2-int8](https://huggingface.co/JustFrederik/nllb-200-1.3B-ct2-int8) | 1.3B | 1.38 GB | int8  | ≈1.3 GB |
+| [JustFrederik/nllb-200-1.3B-ct2-float16](https://huggingface.co/JustFrederik/nllb-200-1.3B-ct2-float16) | 1.3B | 2.74 GB | float16 | ≈1.3 GB |
+| [JustFrederik/nllb-200-distilled-600M-ct2](https://huggingface.co/JustFrederik/nllb-200-distilled-600M-ct2) | 600M | 2.46 GB | float32 | ≈0.6 GB |
+| [JustFrederik/nllb-200-distilled-600M-ct2-float16](https://huggingface.co/JustFrederik/nllb-200-distilled-600M-ct2-float16) | 600M | 1.23 GB | float16 | ≈0.6 GB |
+| [JustFrederik/nllb-200-distilled-600M-ct2-int8](https://huggingface.co/JustFrederik/nllb-200-distilled-600M-ct2-int8) | 600M | 623 MB | int8 | ≈0.6 GB |
+| [JustFrederik/nllb-200-distilled-1.3B-ct2-float16](https://huggingface.co/JustFrederik/nllb-200-distilled-1.3B-ct2-float16) | 1.3B | 2.74 GB | float16 | ≈1.3 GB |
+| [JustFrederik/nllb-200-distilled-1.3B-ct2-int8](https://huggingface.co/JustFrederik/nllb-200-distilled-1.3B-ct2-int8) | 1.3B | 1.38 GB | int8 | ≈1.3 GB |
+| [JustFrederik/nllb-200-distilled-1.3B-ct2](https://huggingface.co/JustFrederik/nllb-200-distilled-1.3B-ct2) | 1.3B | 5.49 GB | float32 | ≈1.3 GB |
+| [JustFrederik/nllb-200-1.3B-ct2](https://huggingface.co/JustFrederik/nllb-200-1.3B-ct2) | 1.3B | 5.49 GB | float32 | ≈1.3 GB |
+| [JustFrederik/nllb-200-3.3B-ct2-float16](https://huggingface.co/JustFrederik/nllb-200-3.3B-ct2-float16) | 3.3B | 6.69 GB | float16 | ≈3.2 GB |
+
+## MT5
+
+mT5 is a multilingual pre-trained Text-to-Text Transformer introduced by Google Research in October 2020. It is a multilingual variant of the T5 model, pre-trained on datasets in 101 languages. Further fine-tuning is required to transform it into a translation model. The paper is titled "`mT5: A Massively Multilingual Pre-trained Text-to-Text Transformer`" ([arXiv:2010.11934](https://arxiv.org/abs/2010.11934)).  
+The 'mt5-zh-ja-en-trimmed' model is finetuned from Google's 'mt5-base' model. This model has a relatively good translation speed, but it only supports three languages: Chinese, Japanese, and English.  
+
+| Name | Parameters | Size | type/quantize | Required VRAM |
+|------|------------|------|---------------|---------------|
+| [mt5-base](https://huggingface.co/google/mt5-base) | N/A | 2.33 GB | float32 | N/A |
+| [K024/mt5-zh-ja-en-trimmed](https://huggingface.co/K024/mt5-zh-ja-en-trimmed) | N/A | 1.32 GB | float32 | ≈1.4 GB |
+| [engmatic-earth/mt5-zh-ja-en-trimmed-fine-tuned-v1](https://huggingface.co/engmatic-earth/mt5-zh-ja-en-trimmed-fine-tuned-v1) | N/A | 1.32 GB | float32 | ≈1.4 GB |
+
+## ALMA
+
+ALMA is a many-to-many LLM-based translation model introduced by Haoran Xu and colleagues in September 2023. It is based on the fine-tuning of a large language model (LLaMA-2). The approach used for this model is referred to as Advanced Language Model-based trAnslator (ALMA). The paper is titled "`A Paradigm Shift in Machine Translation: Boosting Translation Performance of Large Language Models`" ([arXiv:2309.11674](https://arxiv.org/abs/2309.11674)).  
+The official support for ALMA currently includes 10 language directions: English↔German, English↔Czech, English↔Icelandic, English↔Chinese, and English↔Russian. However, the author hints that there might be surprises in other directions, so there are currently no restrictions on the languages that ALMA can be chosen for in the web UI.  
+
+| Name | Parameters | Size | type/quantize | Required VRAM |
+|------|------------|------|---------------|---------------|
+| [haoranxu/ALMA-7B](https://huggingface.co/haoranxu/ALMA-7B) | 7B | 26.95 GB | float32 | N/A |
+| [haoranxu/ALMA-13B](https://huggingface.co/haoranxu/ALMA-13B) | 13B | 52.07 GB | float32 | N/A |
+
+## ALMA-GPTQ
+
+GPTQ is a technique used to quantize the parameters of large language models into integer formats such as int8 or int4. Although the quantization process may lead to a loss in model performance, it significantly reduces both file size and the required VRAM.  
+
+| Name | Parameters | Size | type/quantize | Required VRAM |
+|------|------------|------|---------------|---------------|
+| [TheBloke/ALMA-7B-GPTQ](https://huggingface.co/TheBloke/ALMA-7B-GPTQ) | 7B | 3.9 GB | 4 Bits | ≈4.3 GB |
+| [TheBloke/ALMA-13B-GPTQ](https://huggingface.co/TheBloke/ALMA-13B-GPTQ) | 13B | 7.26 GB | 4 Bits | ≈8.1 |
+
+
+# Options
+
+## Translation - Batch Size
+- transformers: batch_size  
+When the pipeline will use DataLoader (when passing a dataset, on GPU for a Pytorch model), the size of the batch to use, for inference this is not always beneficial.
+- ctranslate2: max_batch_size  
+The maximum batch size.
+
+## Translation - No Repeat Ngram Size
+- transformers: no_repeat_ngram_size  
+Value that will be used by default in the generate method of the model for no_repeat_ngram_size. If set to int > 0, all ngrams of that size can only occur once.
+- ctranslate2: no_repeat_ngram_size  
+Prevent repetitions of ngrams with this size (set 0 to disable).
+
+## Translation - Num Beams
+- transformers: num_beams  
+Number of beams for beam search that will be used by default in the generate method of the model. 1 means no beam search.
+- ctranslate2: beam_size  
+Beam size (1 for greedy search).

requirements-fasterWhisper.txt

@@ -15,4 +15,9 @@ sentencepiece
 intervaltree
 srt
 torch
-https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+
+# Needed by ALMA(GPTQ)
+accelerate
+auto-gptq
+optimum

requirements-whisper.txt

@@ -1,4 +1,5 @@
 git+https://github.com/huggingface/transformers
+ctranslate2>=3.21.0
 git+https://github.com/openai/whisper.git
 ffmpeg-python==0.2.0
 gradio==3.50.2
@@ -13,4 +14,9 @@ sentencepiece
 intervaltree
 srt
 torch
-https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+
+# Needed by ALMA(GPTQ)
+accelerate
+auto-gptq
+optimum

requirements.txt

@@ -15,4 +15,9 @@ sentencepiece
 intervaltree
 srt
 torch
-https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
+
+# Needed by ALMA(GPTQ)
+accelerate
+auto-gptq
+optimum

src/config.py

@@ -5,7 +5,7 @@ from typing import List, Dict, Literal
 
 
 class ModelConfig:
-    def __init__(self, name: str, url: str, path: str = None, type: str = "whisper", tokenizer_url: str = None):
+    def __init__(self, name: str, url: str, path: str = None, type: str = "whisper", tokenizer_url: str = None, revision: str = None):
         """
         Initialize a model configuration.
 
@@ -13,12 +13,17 @@ class ModelConfig:
         url: URL to download the model from
         path: Path to the model file. If not set, the model will be downloaded from the URL.
         type: Type of model. Can be whisper or huggingface.
+        revision: [by transformers] The specific model version to use.
+            It can be a branch name, a tag name, or a commit id, 
+            since we use a git-based system for storing models and other artifacts on huggingface.co, 
+            so revision can be any identifier allowed by git.
         """
         self.name = name
         self.url = url
         self.path = path
         self.type = type
         self.tokenizer_url = tokenizer_url
+        self.revision = revision
 
 VAD_INITIAL_PROMPT_MODE_VALUES=["prepend_all_segments", "prepend_first_segment", "json_prompt_mode"]
 

src/translation/translationLangs.py

@@ -1,4 +1,6 @@
-class Lang():
+from functools import cmp_to_key
+
+class Lang():
     def __init__(self, code: str, *names: str):
         self.code = code
         self.names = names
@@ -292,12 +294,30 @@ def get_lang_whisper_names():
     """Return a list of whisper language names."""
     return list(_TO_LANG_NAME_WHISPER.keys())
 
+def sort_lang_by_whisper_codes(specified_order: list = []):
+    def sort_by_whisper_code(lang: TranslationLang, specified_order: list):
+        return (specified_order.index(lang.whisper.code), lang.whisper.names[0]) if lang.whisper.code in specified_order else (len(specified_order), lang.whisper.names[0])
+
+    def cmp_by_whisper_code(lang1: TranslationLang, lang2: TranslationLang):
+        val1 = sort_by_whisper_code(lang1, specified_order)
+        val2 = sort_by_whisper_code(lang2, specified_order)
+        if val1 > val2:
+            return 1
+        elif val1 == val2:
+            return 0
+        else: return -1
+        
+    sorted_translations = sorted(_TO_LANG_NAME_WHISPER.values(), key=cmp_to_key(cmp_by_whisper_code))
+    return list({name.lower(): None for language in sorted_translations for name in language.whisper.names}.keys())
+
 if __name__ == "__main__":
     # Test lookup
     print("name:Chinese (Traditional)", get_lang_from_nllb_name("Chinese (Traditional)"))
     print("name:moldavian", get_lang_from_m2m100_name("moldavian"))
     print("code:ja", get_lang_from_whisper_code("ja"))
     print("name:English", get_lang_from_nllb_name('English'))
+    print("\n\n")
     
     print(get_lang_m2m100_names(["en", "ja", "zh"]))
-    print(get_lang_nllb_names())
+    print("\n\n")
+    print(sort_lang_by_whisper_codes(["en", "de", "cs", "is", "ru", "zh", "ja"]))

src/translation/translationModel.py

@@ -3,11 +3,9 @@ import warnings
 import huggingface_hub
 import requests
 import torch
-
 import ctranslate2
 import transformers
-
-import re
+import traceback
 
 from typing import Optional
 from src.config import ModelConfig
@@ -85,84 +83,175 @@ class TranslationModel:
             self.load_model()
 
     def load_model(self):
-        print('\n\nLoading model: %s\n\n' % self.modelPath)
-        if "ct2" in self.modelPath:
-            if "nllb" in self.modelPath:
-                self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelConfig.tokenizer_url if self.modelConfig.tokenizer_url is not None and len(self.modelConfig.tokenizer_url) > 0 else self.modelPath, src_lang=self.whisperLang.nllb.code)
-                self.targetPrefix = [self.translationLang.nllb.code]
-            elif "m2m100" in self.modelPath:
-                self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelConfig.tokenizer_url if self.modelConfig.tokenizer_url is not None and len(self.modelConfig.tokenizer_url) > 0 else self.modelPath, src_lang=self.whisperLang.m2m100.code)
-                self.targetPrefix = [self.transTokenizer.lang_code_to_token[self.translationLang.m2m100.code]]
-            self.transModel = ctranslate2.Translator(self.modelPath, compute_type="auto", device=self.device)
-        elif "mt5" in self.modelPath:
-            self.mt5Prefix = self.whisperLang.whisper.code + "2" + self.translationLang.whisper.code + ": "
-            self.transTokenizer = transformers.T5Tokenizer.from_pretrained(self.modelPath, legacy=False) #requires spiece.model
-            self.transModel = transformers.MT5ForConditionalGeneration.from_pretrained(self.modelPath)
-            self.transTranslator = transformers.pipeline('text2text-generation', model=self.transModel, device=self.device, tokenizer=self.transTokenizer)
-        elif "ALMA" in self.modelPath:
-            self.ALMAPrefix = "Translate this from " + self.whisperLang.whisper.code + " to " + self.translationLang.whisper.code + ":" + self.whisperLang.whisper.code + ":"
-            self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelPath, use_fast=True)
-            self.transModel = transformers.AutoModelForCausalLM.from_pretrained(self.modelPath, device_map="auto", trust_remote_code=False, revision="main")
-            self.transTranslator = transformers.pipeline("text-generation", model=self.transModel, tokenizer=self.transTokenizer, batch_size=2, do_sample=True, temperature=0.7, top_p=0.95, top_k=40, repetition_penalty=1.1)
-        else:
-            self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelPath)
-            self.transModel = transformers.AutoModelForSeq2SeqLM.from_pretrained(self.modelPath)
-            if "m2m100" in self.modelPath:
-                self.transTranslator = transformers.pipeline('translation', model=self.transModel, device=self.device, tokenizer=self.transTokenizer, src_lang=self.whisperLang.m2m100.code, tgt_lang=self.translationLang.m2m100.code)
-            else: #NLLB
-                self.transTranslator = transformers.pipeline('translation', model=self.transModel, device=self.device, tokenizer=self.transTokenizer, src_lang=self.whisperLang.nllb.code, tgt_lang=self.translationLang.nllb.code)
+        """
+        [from_pretrained]
+        low_cpu_mem_usage(bool, optional)
+        Tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. This is an experimental feature and a subject to change at any moment.
+        
+        [transformers.AutoTokenizer.from_pretrained]
+        use_fast (bool, optional, defaults to True):
+            Use a fast Rust-based tokenizer if it is supported for a given model. 
+            If a fast tokenizer is not available for a given model, a normal Python-based tokenizer is returned instead.
+        
+        [transformers.AutoModelForCausalLM.from_pretrained]
+        device_map (str or Dict[str, Union[int, str, torch.device], optional):
+            Sent directly as model_kwargs (just a simpler shortcut). When accelerate library is present, 
+            set device_map="auto" to compute the most optimized device_map automatically.
+        revision (str, optional, defaults to "main"):
+            The specific model version to use. It can be a branch name, a tag name, or a commit id, 
+            since we use a git-based system for storing models and other artifacts on huggingface.co, 
+            so revision can be any identifier allowed by git.
+        code_revision (str, optional, defaults to "main")
+            The specific revision to use for the code on the Hub, if the code leaves in a different repository than the rest of the model. 
+            It can be a branch name, a tag name, or a commit id, since we use a git-based system for storing models and other artifacts on huggingface.co, 
+            so revision can be any identifier allowed by git.
+        trust_remote_code (bool, optional, defaults to False):
+            Whether or not to allow for custom models defined on the Hub in their own modeling files. 
+            This option should only be set to True for repositories you trust and in which you have read the code, 
+            as it will execute code present on the Hub on your local machine.
+        
+        [transformers.pipeline "text-generation"]
+        do_sample:
+            if set to True, this parameter enables decoding strategies such as multinomial sampling, 
+            beam-search multinomial sampling, Top-K sampling and Top-p sampling. 
+            All these strategies select the next token from the probability distribution 
+            over the entire vocabulary with various strategy-specific adjustments.
+        temperature (float, optional, defaults to 1.0):
+            The value used to modulate the next token probabilities.
+        top_k (int, optional, defaults to 50):
+            The number of highest probability vocabulary tokens to keep for top-k-filtering.
+        top_p (float, optional, defaults to 1.0):
+            If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation.
+        repetition_penalty (float, optional, defaults to 1.0)
+            The parameter for repetition penalty. 1.0 means no penalty. See this paper for more details.
+        """
+        try:
+            print('\n\nLoading model: %s\n\n' % self.modelPath)
+            if "ct2" in self.modelPath:
+                if any(name in self.modelPath for name in ["nllb", "m2m100"]):
+                    if "nllb" in self.modelPath:
+                        self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelConfig.tokenizer_url if self.modelConfig.tokenizer_url is not None and len(self.modelConfig.tokenizer_url) > 0 else self.modelPath, src_lang=self.whisperLang.nllb.code)
+                        self.targetPrefix = [self.translationLang.nllb.code]
+                    elif "m2m100" in self.modelPath:
+                        self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelConfig.tokenizer_url if self.modelConfig.tokenizer_url is not None and len(self.modelConfig.tokenizer_url) > 0 else self.modelPath, src_lang=self.whisperLang.m2m100.code)
+                        self.targetPrefix = [self.transTokenizer.lang_code_to_token[self.translationLang.m2m100.code]]
+                    self.transModel = ctranslate2.Translator(self.modelPath, compute_type="auto", device=self.device)
+                elif "ALMA" in self.modelPath:
+                    self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelConfig.tokenizer_url if self.modelConfig.tokenizer_url is not None and len(self.modelConfig.tokenizer_url) > 0 else self.modelPath)
+                    self.ALMAPrefix = "Translate this from " + self.whisperLang.whisper.names[0] + " to " + self.translationLang.whisper.names[0] + ":\n" + self.whisperLang.whisper.names[0] + ": "
+                    self.transModel = ctranslate2.Generator(self.modelPath, device=self.device)
+            elif "mt5" in self.modelPath:
+                self.mt5Prefix = self.whisperLang.whisper.code + "2" + self.translationLang.whisper.code + ": "
+                self.transTokenizer = transformers.T5Tokenizer.from_pretrained(self.modelPath, legacy=False) #requires spiece.model
+                self.transModel = transformers.MT5ForConditionalGeneration.from_pretrained(self.modelPath, low_cpu_mem_usage=True)
+                self.transTranslator = transformers.pipeline('text2text-generation', model=self.transModel, device=self.device, tokenizer=self.transTokenizer)
+            elif "ALMA" in self.modelPath:
+                self.ALMAPrefix = "Translate this from " + self.whisperLang.whisper.names[0] + " to " + self.translationLang.whisper.names[0] + ":\n" + self.whisperLang.whisper.names[0] + ": "
+                self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelPath, use_fast=True)
+                self.transModel = transformers.AutoModelForCausalLM.from_pretrained(self.modelPath, device_map="auto", low_cpu_mem_usage=True, trust_remote_code=False, revision=self.modelConfig.revision)
+                self.transTranslator = transformers.pipeline("text-generation", model=self.transModel, tokenizer=self.transTokenizer, do_sample=True, temperature=0.7, top_k=40, top_p=0.95, repetition_penalty=1.1)
+            else:
+                self.transTokenizer = transformers.AutoTokenizer.from_pretrained(self.modelPath)
+                self.transModel = transformers.AutoModelForSeq2SeqLM.from_pretrained(self.modelPath)
+                if "m2m100" in self.modelPath:
+                    self.transTranslator = transformers.pipeline('translation', model=self.transModel, device=self.device, tokenizer=self.transTokenizer, src_lang=self.whisperLang.m2m100.code, tgt_lang=self.translationLang.m2m100.code)
+                else: #NLLB
+                    self.transTranslator = transformers.pipeline('translation', model=self.transModel, device=self.device, tokenizer=self.transTokenizer, src_lang=self.whisperLang.nllb.code, tgt_lang=self.translationLang.nllb.code)
+
+        except Exception as e:
+            print(traceback.format_exc())
+            self.release_vram()
 
     def release_vram(self):
         try:
             if torch.cuda.is_available():
                 if "ct2" not in self.modelPath:
-                    device = torch.device("cpu")
-                    self.transModel.to(device)
+                    try:
+                        device = torch.device("cpu")
+                        self.transModel.to(device)
+                    except Exception as e:
+                        print(traceback.format_exc())
+                        print("\tself.transModel.to cpu, error: " + str(e))
+                    del self.transTranslator
+                del self.transTokenizer
                 del self.transModel
-                torch.cuda.empty_cache()
+                try:
+                    torch.cuda.empty_cache()
+                except Exception as e:
+                    print(traceback.format_exc())
+                    print("\tcuda empty cache, error: " + str(e))
+                import gc
+                gc.collect()
                 print("release vram end.")
         except Exception as e:
             print("Error release vram: " + str(e))
 
 
     def translation(self, text: str, max_length: int = 400):
+        """
+        [ctranslate2]
+        max_batch_size:
+            The maximum batch size. If the number of inputs is greater than max_batch_size, 
+            the inputs are sorted by length and split by chunks of max_batch_size examples 
+            so that the number of padding positions is minimized.
+        no_repeat_ngram_size:
+            Prevent repetitions of ngrams with this size (set 0 to disable).
+        beam_size:
+            Beam size (1 for greedy search).
+        
+        [ctranslate2.Generator.generate_batch]
+        sampling_temperature:
+            Sampling temperature to generate more random samples.
+        sampling_topk:
+            Randomly sample predictions from the top K candidates.
+        sampling_topp:
+            Keep the most probable tokens whose cumulative probability exceeds this value.
+        repetition_penalty:
+            Penalty applied to the score of previously generated tokens (set > 1 to penalize).
+        include_prompt_in_result:
+            Include the start_tokens in the result.
+            If include_prompt_in_result is True (the default), the decoding loop is constrained to generate the start tokens that are then included in the result.
+            If include_prompt_in_result is False, the start tokens are forwarded in the decoder at once to initialize its state (i.e. the KV cache for Transformer models).
+            For variable-length inputs, only the tokens up to the minimum length in the batch are forwarded at once. The remaining tokens are generated in the decoding loop with constrained decoding.
+
+        [transformers.TextGenerationPipeline.__call__]
+        return_full_text (bool, optional, defaults to True):
+            If set to False only added text is returned, otherwise the full text is returned. Only meaningful if return_text is set to True.
+        """
         output = None
         result = None
         try:
             if "ct2" in self.modelPath:
-                source = self.transTokenizer.convert_ids_to_tokens(self.transTokenizer.encode(text))
-                output = self.transModel.translate_batch([source], target_prefix=[self.targetPrefix], max_batch_size=self.batchSize, no_repeat_ngram_size=self.noRepeatNgramSize, beam_size=self.numBeams)
-                target = output[0].hypotheses[0][1:]
-                result = self.transTokenizer.decode(self.transTokenizer.convert_tokens_to_ids(target))
+                if any(name in self.modelPath for name in ["nllb", "m2m100"]):
+                    source = self.transTokenizer.convert_ids_to_tokens(self.transTokenizer.encode(text))
+                    output = self.transModel.translate_batch([source], target_prefix=[self.targetPrefix], max_batch_size=self.batchSize, no_repeat_ngram_size=self.noRepeatNgramSize, beam_size=self.numBeams)
+                    target = output[0].hypotheses[0][1:]
+                    result = self.transTokenizer.decode(self.transTokenizer.convert_tokens_to_ids(target))
+                elif "ALMA" in self.modelPath:
+                    source = self.transTokenizer.convert_ids_to_tokens(self.transTokenizer.encode(self.ALMAPrefix + text + "\n" + self.translationLang.whisper.names[0] + ": "))
+                    output = self.transModel.generate_batch([source], max_length=max_length, max_batch_size=self.batchSize, no_repeat_ngram_size=self.noRepeatNgramSize, beam_size=self.numBeams, sampling_temperature=0.7, sampling_topp=0.9, repetition_penalty=1.1, include_prompt_in_result=False) #, sampling_topk=40
+                    target = output[0]
+                    result = self.transTokenizer.decode(target.sequences_ids[0])
             elif "mt5" in self.modelPath:
                 output = self.transTranslator(self.mt5Prefix + text, max_length=max_length, batch_size=self.batchSize, no_repeat_ngram_size=self.noRepeatNgramSize, num_beams=self.numBeams) #, num_return_sequences=2
                 result = output[0]['generated_text']
             elif "ALMA" in self.modelPath:
-                output = self.transTranslator(self.ALMAPrefix + text + self.translationLang.whisper.code + ":", max_length=max_length, batch_size=self.batchSize, no_repeat_ngram_size=self.noRepeatNgramSize, num_beams=self.numBeams)
+                output = self.transTranslator(self.ALMAPrefix + text + "\n" + self.translationLang.whisper.names[0] + ": ", max_length=max_length, batch_size=self.batchSize, no_repeat_ngram_size=self.noRepeatNgramSize, num_beams=self.numBeams, return_full_text=False)
                 result = output[0]['generated_text']
-                result = re.sub(rf'^(.*{self.translationLang.whisper.code}: )', '', result)  # Remove the prompt from the result
-                result = re.sub(rf'^(Translate this from .* to .*:)', '', result)  # Remove the translation instruction
-                return result.strip()
             else: #M2M100 & NLLB
                 output = self.transTranslator(text, max_length=max_length, batch_size=self.batchSize, no_repeat_ngram_size=self.noRepeatNgramSize, num_beams=self.numBeams)
                 result = output[0]['translation_text']
         except Exception as e:
+            print(traceback.format_exc())
             print("Error translation text: " + str(e))
 
         return result
 
 
-_MODELS = ["distilled-600M", "distilled-1.3B", "1.3B", "3.3B", 
-           "ct2fast-nllb-200-distilled-1.3B-int8_float16", 
-           "ct2fast-nllb-200-3.3B-int8_float16", 
-           "nllb-200-3.3B-ct2-float16", "nllb-200-1.3B-ct2", "nllb-200-1.3B-ct2-int8", "nllb-200-1.3B-ct2-float16", 
-           "nllb-200-distilled-1.3B-ct2", "nllb-200-distilled-1.3B-ct2-int8", "nllb-200-distilled-1.3B-ct2-float16", 
-           "nllb-200-distilled-600M-ct2", "nllb-200-distilled-600M-ct2-int8", "nllb-200-distilled-600M-ct2-float16",
-           "m2m100_1.2B-ct2", "m2m100_418M-ct2", "m2m100-12B-ct2", 
-           "m2m100_1.2B", "m2m100_418M",
-           "mt5-zh-ja-en-trimmed",
-           "mt5-zh-ja-en-trimmed-fine-tuned-v1",
-           "ALMA-13B-GPTQ"]
+_MODELS = ["nllb-200", 
+           "m2m100",
+           "mt5",
+           "ALMA"]
 
 def check_model_name(name):
     return any(allowed_name in name for allowed_name in _MODELS)
@@ -230,6 +319,9 @@ def download_model(
         "allow_patterns": allowPatterns,
         #"tqdm_class": disabled_tqdm,
     }
+    
+    if modelConfig.revision is not None:
+        kwargs["revision"] = modelConfig.revision
 
     if outputDir is not None:
         kwargs["local_dir"] = outputDir

src/utils.py

@@ -130,7 +130,7 @@ def write_srt_original(transcript: Iterator[dict], file: TextIO,
             flush=True,
         )
         
-        if original is not None: print(f"{original}\n",
+        if original is not None: print(f"{original}",
             file=file,
             flush=True)