Update src/translate/Translate.py

src/translate/Translate.py

@@ -2,6 +2,8 @@ from nltk.tokenize import sent_tokenize
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import src.exception.Exception.Exception as ExceptionCustom
+# Use a pipeline as a high-level helper
+from transformers import pipeline
 
 METHOD = "TRANSLATE"
 
@@ -49,4 +51,11 @@ def paraphraseTranslateMethod(requestValue: str, model: str):
             result = tokenizerENGROM.batch_decode(output, skip_special_tokens=True)[0]
         result_value.append(result)
 
-    return " ".join(result_value).strip(), model
+    return " ".join(result_value).strip(), model
+
+def gemma(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-better'):  
+    pipe = pipeline("text-generation", model="Gargaz/gemma-2b-romanian-better")
+    messages = [
+        {"role": "user", "content": f"Translate the following text to Romanian using a formal tone and provide only translation: {requestValue}"},
+    ]
+    return pipe(messages)