separated the translator that uses the HF endpoint

Dockerfile

@@ -1,16 +1,4 @@
 FROM python:3.12-slim
-## Set up a new user named "user" with user ID 1000
-#RUN useradd -m -u 1000 user
-#
-## Switch to the "user" user
-#USER user
-#
-## Set home to the user's home directory
-#ENV HOME=/home/user \
-#	PATH=/home/user/.local/bin:$PATH
-
-## Set the working directory to the user's home directory
-#WORKDIR $HOME/app
 
 COPY src ./src
 COPY scripts ./scripts

gradio_app.py

@@ -1,6 +1,5 @@
 import gradio as gr
-from src.translate_any_doc import translate_document
-from src.salamandraTA7b_translator import SalamandraTA7bTranslator
+from src.salamandraTA7b_translator_HF import SalamandraTA7bTranslatorHF
 from src.aligner import Aligner
 import os
 
@@ -12,9 +11,9 @@ def upload_file(filepath, source_lang, target_lang, user_token):
     hf_token = os.environ.get('HF_TOKEN')
     if user_token:
         hf_token = user_token
-    translator = SalamandraTA7bTranslator(hf_token)
+    translator = SalamandraTA7bTranslatorHF(hf_token)
     aligner = Aligner(config_folder, source_lang, target_lang, temp_folder)
-    for status, translated_file_name in translate_document(filepath, source_lang, target_lang, translator, aligner):
+    for status, translated_file_name in translator.translate_document(filepath, source_lang, target_lang, aligner):
         if translated_file_name:  # finished
             yield [gr.UploadButton(visible=False),
                    gr.DownloadButton(label=f"Download {translated_file_name}", value=translated_file_name,

requirements.txt

@@ -1,8 +1,12 @@
-iso-639~=0.4.5
-protobuf~=6.30.2
 requests~=2.32.3
 tqdm~=4.67.1
 gradio~=5.25.1
+spacy~=3.8.6
 gradio_client~=1.8.0
-setuptools~=80.0.0
-spacy~=3.8.6
+pandas~=2.3.3
+beautifulsoup4~=4.14.2
+gradio_client~=1.8.0
+transformers~=4.57.1
+torch~=2.8.0
+huggingface-hub~=0.36.0
+vllm~=0.11.0

src/{translate_any_doc.py → salamandraTA7b_translator_HF.py}

@@ -1,6 +1,6 @@
 import shutil
 import string
-import sys
+from iso639 import languages
 import time
 import os
 from itertools import groupby
@@ -23,15 +23,154 @@ spacy_nlp = spacy.load("xx_ent_wiki_sm")
 if "sentencizer" not in spacy_nlp.pipe_names:
     spacy_nlp.add_pipe("sentencizer")
 
-
 import unicodedata
 
+
+class SalamandraTA7bTranslatorHF:
+    def __init__(self, hf_token):
+        from gradio_client import Client
+        self.client = Client("BSC-LT/SalamandraTA-7B-Demo", hf_token=hf_token)
+
+    def translate(self, text, source_lang, target_lang):
+        if not text:
+            return ""
+
+        # 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]
+
+    def translate_document(self, input_file: str, source_lang: str, target_lang: str,
+                           aligner: Aligner,
+                           temp_folder: str = "tmp",
+                           tikal_folder: str = "okapi-apps_gtk2-linux-x86_64_1.47.0") -> (str, str):
+        input_filename = input_file.split("/")[-1]
+        os.makedirs(temp_folder, exist_ok=True)
+
+        # copy the original file to the temporal folder to avoid common issues with tikal
+        temp_input_file = os.path.join(temp_folder, input_filename)
+        shutil.copy(input_file, temp_input_file)
+
+        original_xliff_file = os.path.join(temp_folder, input_filename + ".xlf")
+        plain_text_file = doc_to_plain_text(temp_input_file, source_lang, target_lang, tikal_folder,
+                                            original_xliff_file)
+
+        # get paragraphs with runs
+        paragraphs_with_runs = [get_runs_from_paragraph(line.strip(), idx) for idx, line in
+                                enumerate(open(plain_text_file).readlines())]
+
+        # translate using plaintext file
+        original_tokenized_sentences_with_style = []
+        original_spacing = []
+        for run in paragraphs_with_runs:
+            tokens, spaces = tokenize_with_runs(run)
+            original_tokenized_sentences_with_style += tokens
+            original_spacing += spaces
+
+        translated_sentences = []
+        yield "Translating 0%...", None
+        total = len(original_tokenized_sentences_with_style)
+        pbar = tqdm.tqdm(desc="Translating paragraphs...", total=total)
+
+        for i, (sentence, spacing) in enumerate(zip(original_tokenized_sentences_with_style, original_spacing)):
+            text = Doc(spacy_nlp.vocab, words=[token["text"] for token in sentence], spaces=spacing).text
+
+            while True:
+                try:
+                    translated_sentences.append(self.translate(text, source_lang, target_lang))
+                    break
+                except AppError as e:
+                    print(e)
+
+            pbar.update(1)
+            percent_complete = int(((i + 1) / total) * 100)
+            yield f"Translating {percent_complete}%...", None
+
+        # time to align the translation with the original
+        print("Generating alignments...")
+        yield "Aligning...", None
+        start_time = time.time()
+        translated_sentences_with_style, translated_sentences_spacing = generate_alignments(
+            original_tokenized_sentences_with_style,
+            translated_sentences, aligner,
+            temp_folder)
+        print(f"Finished alignments in {time.time() - start_time} seconds")
+
+        # since we tokenized these sentences independently, the spacing information does not contain spaces after punctuation
+        # at the end of the sentence (there's no space at the end of a sentence that ends with ".", unless there's a sentence
+        # right after
+        for sentence, sentence_spaces in zip(translated_sentences_with_style, translated_sentences_spacing):
+            if sentence[-1]["text"] in string.punctuation:
+                sentence_spaces[-1] = True
+
+        # flatten the sentences into a list of tokens
+        translated_tokens_with_style = [item for sublist in translated_sentences_with_style for item in sublist]
+        tokens_spaces = [item for sublist in translated_sentences_spacing for item in sublist]
+
+        # group the tokens by style/run
+        translated_runs_with_style = group_by_style(translated_tokens_with_style, tokens_spaces)
+
+        # group the runs by original paragraph
+        translated_paragraphs_with_style = {key: [{'id': None, 'paragraph_index': key, 'text': ""}] for key in
+                                            range(len(paragraphs_with_runs))}
+
+        for item in translated_runs_with_style:
+            # first item in the paragraph, remove starting blank space we introduced in group_by_style(), where we
+            # didn't know where paragraphs started and ended
+            if not translated_paragraphs_with_style[item['paragraph_index']][0]["text"]:
+                first_item_in_paragraph = item.copy()
+                first_item_in_paragraph["text"] = first_item_in_paragraph["text"].lstrip(" ")
+                translated_paragraphs_with_style[item['paragraph_index']] = []
+                translated_paragraphs_with_style[item['paragraph_index']].append(first_item_in_paragraph)
+            else:
+                translated_paragraphs_with_style[item['paragraph_index']].append(item)
+
+        # save to new plain text file
+        translated_moses_file = os.path.join(original_xliff_file + f".{target_lang}")
+        runs_to_plain_text(translated_paragraphs_with_style, translated_moses_file)
+
+        # put the translations into the xlf
+        tikal_moses_to_xliff_command = [os.path.join(tikal_folder, "tikal.sh"), "-lm", original_xliff_file, "-sl",
+                                        source_lang, "-tl", target_lang, "-from", translated_moses_file, "-totrg",
+                                        "-noalttrans", "-to", original_xliff_file]
+        Popen(tikal_moses_to_xliff_command).wait()
+
+        # any tags that are still <g> have not been paired between original and translated texts by tikal so we remove
+        # them. This may happen if a word in the original language has been split in more that one words that have other
+        # words in between, or an error in fastalign
+        text = open(original_xliff_file).read()
+        result = re.sub(r'<g id="\d+">(.*?)</g>', r'\1', text)
+        open(original_xliff_file, "w").write(result)
+
+        # merge into a docx again
+        tikal_merge_doc_command = [os.path.join(tikal_folder, "tikal.sh"), "-m", original_xliff_file]
+        final_process = Popen(tikal_merge_doc_command, stdout=PIPE, stderr=PIPE)
+        stdout, stderr = final_process.communicate()
+        final_process.wait()
+
+        # get the path to the output file
+        output = stdout.decode('utf-8')
+        translated_file_path = re.search(r'(?<=Output:\s)(.*)', output)[0]
+
+        print(f"Saved file in {translated_file_path}")
+        yield "", translated_file_path
+
+
 def remove_invisible(text):
     return ''.join(
         c for c in text
         if not unicodedata.category(c) in ['Zs', 'Cc', 'Cf']
     )
 
+
 def get_leading_invisible(text):
     i = 0
     while i < len(text):
@@ -42,6 +181,7 @@ def get_leading_invisible(text):
             break
     return text[:i]
 
+
 def doc_to_plain_text(input_file: str, source_lang: str, target_lang: str, tikal_folder: str,
                       original_xliff_file_path: str) -> str:
     """
@@ -118,7 +258,7 @@ def get_runs_from_paragraph(paragraph: str, paragraph_index: int) -> list[dict[s
             # Self-closing tag like <x id="1"/>
             if tag_id is None:
                 tag_id = -1
-                #raise ValueError(f"Self-closing tag <{tag_name}/> missing id")
+                # raise ValueError(f"Self-closing tag <{tag_name}/> missing id")
             runs.append({
                 "text": "",
                 "id": [f"{tag_name}_{tag_id}"],
@@ -128,7 +268,7 @@ def get_runs_from_paragraph(paragraph: str, paragraph_index: int) -> list[dict[s
             # Opening tag <tag id="...">
             if tag_id is None:
                 tag_id = -1
-                #raise ValueError(f"Opening tag <{tag_name}> missing id")
+                # raise ValueError(f"Opening tag <{tag_name}> missing id")
             tag_stack.append(f"{tag_name}_{tag_id}")
 
         pos = end
@@ -419,120 +559,3 @@ def runs_to_plain_text(paragraphs_with_style: dict[int, list[dict[str, str, str]
                     out_file.write("".join(run["text"]))
 
             out_file.write("\n")
-
-
-def translate_document(input_file: str, source_lang: str, target_lang: str,
-                       translator,
-                       aligner: Aligner,
-                       temp_folder: str = "tmp",
-                       tikal_folder: str = "okapi-apps_gtk2-linux-x86_64_1.47.0", with_format: bool = True) -> (str,
-                                                                                                                str):
-    input_filename = input_file.split("/")[-1]
-    os.makedirs(temp_folder, exist_ok=True)
-
-    # copy the original file to the temporal folder to avoid common issues with tikal
-    temp_input_file = os.path.join(temp_folder, input_filename)
-    shutil.copy(input_file, temp_input_file)
-
-    original_xliff_file = os.path.join(temp_folder, input_filename + ".xlf")
-    plain_text_file = doc_to_plain_text(temp_input_file, source_lang, target_lang, tikal_folder, original_xliff_file)
-
-    # get paragraphs with runs
-    paragraphs_with_runs = [get_runs_from_paragraph(line.strip(), idx) for idx, line in
-                            enumerate(open(plain_text_file).readlines())]
-
-    # translate using plaintext file
-    original_tokenized_sentences_with_style = []
-    original_spacing = []
-    for run in paragraphs_with_runs:
-        tokens, spaces = tokenize_with_runs(run)
-        original_tokenized_sentences_with_style += tokens
-        original_spacing += spaces
-
-    translated_sentences = []
-    yield "Translating 0%...", None
-    total = len(original_tokenized_sentences_with_style)
-    pbar = tqdm.tqdm(desc="Translating paragraphs...", total=total)
-
-    for i, (sentence, spacing) in enumerate(zip(original_tokenized_sentences_with_style, original_spacing)):
-        text = Doc(spacy_nlp.vocab, words=[token["text"] for token in sentence], spaces=spacing).text
-
-        while True:
-            try:
-                translated_sentences.append(translator.translate(text, source_lang, target_lang))
-                break
-            except AppError as e:
-                print(e)
-
-        pbar.update(1)
-        percent_complete = int(((i + 1) / total) * 100)
-        yield f"Translating {percent_complete}%...", None
-
-    # time to align the translation with the original
-    print("Generating alignments...")
-    yield "Aligning...", None
-    start_time = time.time()
-    translated_sentences_with_style, translated_sentences_spacing = generate_alignments(
-        original_tokenized_sentences_with_style,
-        translated_sentences, aligner,
-        temp_folder)
-    print(f"Finished alignments in {time.time() - start_time} seconds")
-
-    # since we tokenized these sentences independently, the spacing information does not contain spaces after punctuation
-    # at the end of the sentence (there's no space at the end of a sentence that ends with ".", unless there's a sentence
-    # right after
-    for sentence, sentence_spaces in zip(translated_sentences_with_style, translated_sentences_spacing):
-        if sentence[-1]["text"] in string.punctuation:
-            sentence_spaces[-1] = True
-
-    # flatten the sentences into a list of tokens
-    translated_tokens_with_style = [item for sublist in translated_sentences_with_style for item in sublist]
-    tokens_spaces = [item for sublist in translated_sentences_spacing for item in sublist]
-
-    # group the tokens by style/run
-    translated_runs_with_style = group_by_style(translated_tokens_with_style, tokens_spaces)
-
-    # group the runs by original paragraph
-    translated_paragraphs_with_style = {key: [{'id': None, 'paragraph_index': key, 'text': ""}] for key in
-                                        range(len(paragraphs_with_runs))}
-
-    for item in translated_runs_with_style:
-        # first item in the paragraph, remove starting blank space we introduced in group_by_style(), where we
-        # didn't know where paragraphs started and ended
-        if not translated_paragraphs_with_style[item['paragraph_index']][0]["text"]:
-            first_item_in_paragraph = item.copy()
-            first_item_in_paragraph["text"] = first_item_in_paragraph["text"].lstrip(" ")
-            translated_paragraphs_with_style[item['paragraph_index']] = []
-            translated_paragraphs_with_style[item['paragraph_index']].append(first_item_in_paragraph)
-        else:
-            translated_paragraphs_with_style[item['paragraph_index']].append(item)
-
-    # save to new plain text file
-    translated_moses_file = os.path.join(original_xliff_file + f".{target_lang}")
-    runs_to_plain_text(translated_paragraphs_with_style, translated_moses_file)
-
-    # put the translations into the xlf
-    tikal_moses_to_xliff_command = [os.path.join(tikal_folder, "tikal.sh"), "-lm", original_xliff_file, "-sl",
-                                    source_lang, "-tl", target_lang, "-from", translated_moses_file, "-totrg",
-                                    "-noalttrans", "-to", original_xliff_file]
-    Popen(tikal_moses_to_xliff_command).wait()
-
-    # any tags that are still <g> have not been paired between original and translated texts by tikal so we remove
-    # them. This may happen if a word in the original language has been split in more that one words that have other
-    # words in between, or an error in fastalign
-    text = open(original_xliff_file).read()
-    result = re.sub(r'<g id="\d+">(.*?)</g>', r'\1', text)
-    open(original_xliff_file, "w").write(result)
-
-    # merge into a docx again
-    tikal_merge_doc_command = [os.path.join(tikal_folder, "tikal.sh"), "-m", original_xliff_file]
-    final_process = Popen(tikal_merge_doc_command, stdout=PIPE, stderr=PIPE)
-    stdout, stderr = final_process.communicate()
-    final_process.wait()
-
-    # get the path to the output file
-    output = stdout.decode('utf-8')
-    translated_file_path = re.search(r'(?<=Output:\s)(.*)', output)[0]
-
-    print(f"Saved file in {translated_file_path}")
-    yield "", translated_file_path