Update src/translate/Translate.py

src/translate/Translate.py

@@ -7,17 +7,6 @@ from transformers import pipeline
 
 METHOD = "TRANSLATE"
 
-# Load models and tokenizers
-tokenizerROMENG = AutoTokenizer.from_pretrained("BlackKakapo/opus-mt-ro-en")
-modelROMENG = AutoModelForSeq2SeqLM.from_pretrained("BlackKakapo/opus-mt-ro-en")
-
-tokenizerENGROM = AutoTokenizer.from_pretrained("BlackKakapo/opus-mt-en-ro")
-modelENGROM = AutoModelForSeq2SeqLM.from_pretrained("BlackKakapo/opus-mt-en-ro")
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-modelROMENG.to(device)
-modelENGROM.to(device)
-
 def paraphraseTranslateMethod(requestValue: str, model: str):
     exception = ExceptionCustom.checkForException(requestValue, METHOD)
     if exception:
@@ -25,9 +14,13 @@ def paraphraseTranslateMethod(requestValue: str, model: str):
 
     tokenized_sent_list = sent_tokenize(requestValue)
     result_value = []
-
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
     for SENTENCE in tokenized_sent_list:
         if model == 'roen':
+            tokenizerROMENG = AutoTokenizer.from_pretrained("BlackKakapo/opus-mt-ro-en")
+            modelROMENG = AutoModelForSeq2SeqLM.from_pretrained("BlackKakapo/opus-mt-ro-en")
+            modelROMENG.to(device)
             input_ids = tokenizerROMENG(SENTENCE, return_tensors='pt').to(device)
             output = modelROMENG.generate(
                 input_ids=input_ids.input_ids,
@@ -39,6 +32,9 @@ def paraphraseTranslateMethod(requestValue: str, model: str):
             )
             result = tokenizerROMENG.batch_decode(output, skip_special_tokens=True)[0]
         else:
+            tokenizerENGROM = AutoTokenizer.from_pretrained("BlackKakapo/opus-mt-en-ro")
+            modelENGROM = AutoModelForSeq2SeqLM.from_pretrained("BlackKakapo/opus-mt-en-ro")
+            modelENGROM.to(device)
             input_ids = tokenizerENGROM(SENTENCE, return_tensors='pt').to(device)
             output = modelENGROM.generate(
                 input_ids=input_ids.input_ids,