rename

main.py

@@ -19,10 +19,10 @@ class TranslationRequest(BaseModel):
     target_lang: str
 
 
-def load_model(pretrained_model: str = "facebook/m2m100_1.2B", cache_dir: str = "models/"):
+def load_model(model: str = "facebook/m2m100_1.2B", cache_dir: str = "models/"):
     model_dir = os.path.join(os.getcwd(), cache_dir)
-    tokenizer = M2M100Tokenizer.from_pretrained(pretrained_model, cache_dir=model_dir)
-    model = M2M100ForConditionalGeneration.from_pretrained(pretrained_model, cache_dir=model_dir).to(device)
+    tokenizer = M2M100Tokenizer.from_pretrained(model, cache_dir=model_dir)
+    model = M2M100ForConditionalGeneration.from_pretrained(model, cache_dir=model_dir).to(device)
     model.eval()
     return tokenizer, model
 # aparentemente temos um problema ao carregar o modelo então vou tentar carregar no start da aplicação para não dar time-out na request