added salamandraTA translation, update requirements

main.py

@@ -17,6 +17,46 @@ from subprocess import Popen, PIPE
 from itertools import groupby
 import fileinput
 
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from iso639 import languages
+import tqdm
+
+
+class Translator():
+    def __init__(self, model_path, source_lang, target_lang):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            device_map="auto",
+            torch_dtype=torch.bfloat16
+        )
+
+        self.prompt_f = lambda x: (f"Translate the following text from {source_lang} into "
+                                   f"{target_lang}.\n{source_lang}: {x} \n{target_lang}:")
+
+    def translate(self, text):
+        message = [{"role": "user", "content": self.prompt_f(text)}]
+        date_string = datetime.today().strftime('%Y-%m-%d')
+
+        prompt = self.tokenizer.apply_chat_template(
+            message,
+            tokenize=False,
+            add_generation_prompt=True,
+            date_string=date_string
+        )
+
+        inputs = self.tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+        input_length = inputs.shape[1]
+        outputs = self.model.generate(input_ids=inputs.to(self.model.device),
+                                      max_new_tokens=400,
+                                      early_stopping=True,
+                                      num_beams=5)
+
+        return self.tokenizer.decode(outputs[0, input_length:], skip_special_tokens=True)
+
 
 # Class to align original and translated sentences
 # based on https://github.com/mtuoc/MTUOC-server/blob/main/GetWordAlignments_fast_align.py
@@ -235,12 +275,6 @@ def generate_alignments(original_paragraphs_with_runs, translated_paragraphs, al
     return translated_sentences_with_style
 
 
-# TODO
-def translate_paragraph(paragraph_text):
-    translated_paragraph = ""
-    return translated_paragraphs
-
-
 # group contiguous elements with the same boolean values
 def group_by_style(values, detokenizer):
     groups = []
@@ -316,12 +350,17 @@ if __name__ == "__main__":
 
     detokenizer = TreebankWordDetokenizer()
 
+    translator = Translator("BSC-LT/salamandraTA-7b-instruct", languages.get(alpha2=source_lang).name,
+                            languages.get(alpha2=target_lang).name)
+
     # translate each paragraph
     translated_paragraphs = []
-    for paragraph in paragraphs_with_runs:
+    for paragraph in tqdm.tqdm(paragraphs_with_runs, desc="Translating paragraphs..."):
         paragraph_text = detokenizer.detokenize([run["text"] for run in paragraph])
-        translated_paragraphs.append(translate_paragraph(paragraph_text))
-    
+        translated_paragraphs.append(translator.translate(paragraph_text))
+
+    print(translated_paragraphs)
+
     out_doc = Document()
 
     processed_original_paragraphs_with_runs = [preprocess_runs(runs) for runs in paragraphs_with_runs]

requirements.txt

@@ -1,2 +1,7 @@
 nltk~=3.9.1
-python-docx~=1.1.2
+python-docx~=1.1.2
+torch~=2.6.0
+transformers~=4.51.2
+iso-639~=0.4.5
+protobuf~=6.30.2
+sentencepiece~=0.2.0