[Update] Add compute type management for CTranslate2 and Whisper models
This commit is contained in:
@@ -36,14 +36,15 @@ class Translator():
|
||||
result = False
|
||||
return result
|
||||
|
||||
def changeCTranslate2Model(self, path, model_type, device="cpu", device_index=0):
|
||||
def changeCTranslate2Model(self, path, model_type, device="cpu", device_index=0, compute_type="auto"):
|
||||
self.is_loaded_ctranslate2_model = False
|
||||
directory_name = ctranslate2_weights[model_type]["directory_name"]
|
||||
tokenizer = ctranslate2_weights[model_type]["tokenizer"]
|
||||
weight_path = os_path.join(path, "weights", "ctranslate2", directory_name)
|
||||
tokenizer_path = os_path.join(path, "weights", "ctranslate2", directory_name, "tokenizer")
|
||||
|
||||
compute_type = getBestComputeType(device, device_index)
|
||||
if compute_type == "auto":
|
||||
compute_type = getBestComputeType(device, device_index)
|
||||
self.ctranslate2_translator = ctranslate2.Translator(
|
||||
weight_path,
|
||||
device=device,
|
||||
|
||||
Reference in New Issue
Block a user