v0.3.4b4

wasertech · Jul 22, 2023 · 4621e83 · 4621e83
1 parent a419986
commit 4621e83
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/translator/__init__.py b/translator/__init__.py
@@ -1,6 +1,6 @@
 from translator.translate import Translator
 from translator.language import get_nllb_lang
 
-__version__ = "0.3.4b3"
+__version__ = "0.3.4b4"
 
 LANGS = get_nllb_lang()
diff --git a/translator/translate.py b/translator/translate.py
@@ -20,7 +20,7 @@ def __init__(self, source_language, target_language, max_length=500, model_id="f
         self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
         self.logger.debug(f"{self.device}")
         self.logger.debug("Loading model...")
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id, model_max_length=max_length)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id, max_length=max_length)
         self.logger.debug("Loading tokenizer...")
         self.tokenizer = AutoTokenizer.from_pretrained(model_id, model_max_length=max_length)
         self.logger.debug("Setting up translation pipeline...")