Updated readme and added salamandraTA7b translator class

readme.md

@@ -1,6 +1,6 @@
 # document_translator
 
-Project to translate files (for now .docx) using BSC's models while keeping the formatting and style of the original file.
+Project to translate files using BSC's models while keeping the formatting and style of the original file.
 
 ## Requirements
 ### python 3.12
@@ -16,3 +16,12 @@ I took the 4 files (ca-en.params, ca-en.err, en-ca.params and en-ca.err) from ht
 ### python requirements
 
     pip install -r requirements.txt
+
+### mtuoc_aina_translator
+
+To use this class you also need to be running MTUOC's translation server with the proper translation models. There's also no 
+need to use fastalign on that side since the current project already runs it. 
+
+### salamandrata7b_translator
+
+Class that uses huggingface's demo.

src/mtuoc_aina_translator.py

@@ -1,20 +1,19 @@
 import requests
 import json
-from nltk.tokenize import sent_tokenize
+
 
 class MTUOCAinaTranslator:
     def __init__(self, ip: str, port: str):
         self.ip = ip
         self.port = port
 
-    def translate(self, text):
-        stuff = sent_tokenize(text)
-
+    def translate(self, text, source_lang=None, target_lang=None):
         myobj = {
             'id': '1',
             'src': text,
         }
-        url = 'http://' + self.ip + ':' + self.port + '/translate'
+        url = f'http://{self.ip}:{self.port}/translate'
+        #url = 'http://' + self.ip + ':' + self.port + '/translate'
         x = requests.post(url, json=myobj)
         json_response = json.loads(x.text)
         return json_response['tgt']

src/salamandraTA7b_translator.py

@@ -0,0 +1,22 @@
+from gradio_client import Client
+from iso639 import languages
+
+HF_TOKEN = "YOUR-HF-TOKEN-HERE"
+
+class SalamandraTA7bTranslator:
+    def __init__(self):
+        self.client = Client("BSC-LT/SalamandraTA-7B-Demo", hf_token=HF_TOKEN)
+
+    def translate(self, text, source_lang, target_lang):
+        # we assume that they are specifying the language by code so we need to convert it to name
+        lang1 = languages.get(alpha2=source_lang).name
+        lang2 = languages.get(alpha2=target_lang).name
+        result = self.client.predict(
+            task="Translation",
+            source=lang1,
+            target=lang2,
+            input_text=text,
+            mt_text=None,
+            api_name="/generate_output"
+        )
+        return result[0]

src/translate_any_doc.py

@@ -271,7 +271,7 @@ def translate_document(input_file: str, source_lang: str, target_lang: str,
     translated_paragraphs = []
     for paragraph in tqdm.tqdm(paragraphs_with_runs, desc="Translating paragraphs..."):
         paragraph_text = detokenizer.detokenize([run["text"] for run in paragraph])
-        translated_paragraphs.append(translator.translate(paragraph_text))
+        translated_paragraphs.append(translator.translate(paragraph_text, source_lang, target_lang))
 
     # time to align the translation with the original
     print("Generating alignments...")