Spaces:
Sleeping
Sleeping
| import os | |
| import string | |
| from docx import Document | |
| from docx.text.hyperlink import Hyperlink | |
| from docx.text.run import Run | |
| import nltk | |
| nltk.download('punkt') | |
| nltk.download('punkt_tab') | |
| from nltk.tokenize import sent_tokenize, word_tokenize | |
| from nltk.tokenize.treebank import TreebankWordDetokenizer | |
| from subprocess import Popen, PIPE | |
| from itertools import groupby | |
| import fileinput | |
| # Class to align original and translated sentences | |
| # based on https://github.com/mtuoc/MTUOC-server/blob/main/GetWordAlignments_fast_align.py | |
| class Aligner(): | |
| def __init__(self, config_folder, source_lang, target_lang, temp_folder): | |
| forward_params_path = os.path.join(config_folder, f"{source_lang}-{target_lang}.params") | |
| reverse_params_path = os.path.join(config_folder, f"{target_lang}-{source_lang}.params") | |
| fwd_T, fwd_m = self.__read_err(os.path.join(config_folder, f"{source_lang}-{target_lang}.err")) | |
| rev_T, rev_m = self.__read_err(os.path.join(config_folder, f"{target_lang}-{source_lang}.err")) | |
| self.forward_alignment_file_path = os.path.join(temp_folder, "forward.align") | |
| self.reverse_alignment_file_path = os.path.join(temp_folder, "reverse.align") | |
| self.forward_command = lambda \ | |
| x: f'./fast_align -i {x} -d -T {fwd_T} -m {fwd_m} -f {forward_params_path} > {self.forward_alignment_file_path}' | |
| self.reverse_command = lambda \ | |
| x: f'./fast_align -i {x} -d -T {rev_T} -m {rev_m} -f {reverse_params_path} -r > {self.reverse_alignment_file_path}' | |
| self.symmetric_command = f'./atools -i {self.forward_alignment_file_path} -j {self.reverse_alignment_file_path} -c grow-diag-final-and' | |
| def __simplify_alignment_file(self, file): | |
| with fileinput.FileInput(file, inplace=True, backup='.bak') as f: | |
| for line in f: | |
| print(line.split('|||')[2].strip()) | |
| def __read_err(self, err): | |
| (T, m) = ('', '') | |
| for line in open(err): | |
| # expected target length = source length * N | |
| if 'expected target length' in line: | |
| m = line.split()[-1] | |
| # final tension: N | |
| elif 'final tension' in line: | |
| T = line.split()[-1] | |
| return T, m | |
| def align(self, file): | |
| # generate forward alignment | |
| process = Popen(self.forward_command(file), shell=True) | |
| process.wait() | |
| # generate reverse alignment | |
| process = Popen(self.reverse_command(file), shell=True) | |
| process.wait() | |
| # for some reason the output file contains more information than needed, remove it | |
| self.__simplify_alignment_file(self.forward_alignment_file_path) | |
| self.__simplify_alignment_file(self.reverse_alignment_file_path) | |
| # generate symmetrical alignment | |
| process = Popen(self.symmetric_command, shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE) | |
| process.wait() | |
| # get final alignments and format them | |
| alignments_str = process.communicate()[0].decode('utf-8') | |
| alignments = [] | |
| for line in alignments_str.splitlines(): | |
| alignments.append([(int(i), int(j)) for i, j in [pair.split("-") for pair in line.strip("\n").split(" ")]]) | |
| return alignments | |
| # Function to extract paragraphs with their runs | |
| def extract_paragraphs_with_runs(doc): | |
| paragraphs_with_runs = [] | |
| for paragraph in doc.paragraphs: | |
| runs = [] | |
| for item in paragraph.iter_inner_content(): | |
| if isinstance(item, Run): | |
| runs.append({ | |
| 'text': item.text, | |
| 'bold': item.bold, | |
| 'italic': item.italic, | |
| 'underline': item.underline, | |
| 'font_name': item.font.name, | |
| 'font_size': item.font.size, | |
| 'font_color': item.font.color.rgb | |
| }) | |
| elif isinstance(item, Hyperlink): | |
| runs.append({ | |
| 'text': item.runs[0].text, | |
| 'bold': item.runs[0].bold, | |
| 'italic': item.runs[0].italic, | |
| 'underline': item.runs[0].underline, | |
| 'font_name': item.runs[0].font.name, | |
| 'font_size': item.runs[0].font.size, | |
| 'font_color': item.runs[0].font.color.rgb | |
| }) | |
| paragraphs_with_runs.append(runs) | |
| return paragraphs_with_runs | |
| def tokenize_paragraph_with_runs2(runs_in_paragraph): | |
| text_paragraph = " ".join(run["text"] for run in runs_in_paragraph) | |
| sentences = sent_tokenize(text_paragraph) | |
| tokenized_sentences = [word_tokenize(sentence) for sentence in sentences] | |
| tokenized_sentences_with_style = [] | |
| for tokenized_sentence in tokenized_sentences: | |
| tokenized_sentence_with_style = [] | |
| token_idx = 0 | |
| for run in runs_in_paragraph: | |
| text_in_run = run["text"].strip() | |
| if text_in_run == tokenized_sentence[token_idx]: | |
| new_run = run.copy() | |
| new_run["text"] = text_in_run | |
| tokenized_sentence_with_style.append(new_run) | |
| token_idx += 1 | |
| if token_idx >= len(tokenized_sentence): | |
| break | |
| elif len(text_in_run) > len(tokenized_sentence[token_idx]): | |
| if text_in_run.startswith(tokenized_sentence[token_idx]): | |
| for token in word_tokenize(text_in_run): | |
| if token == tokenized_sentence[token_idx]: | |
| new_run = run.copy() | |
| new_run["text"] = token | |
| tokenized_sentence_with_style.append(new_run) | |
| token_idx += 1 | |
| else: | |
| raise "oops" | |
| tokenized_sentences_with_style.append(tokenized_sentence_with_style) | |
| return tokenized_sentences_with_style | |
| def tokenize_paragraph_with_runs(runs_in_paragraph, detokenizer): | |
| text_paragraph = detokenizer.detokenize([run["text"] for run in runs_in_paragraph]) | |
| sentences = sent_tokenize(text_paragraph) | |
| tokenized_sentences = [word_tokenize(sentence) for sentence in sentences] | |
| tokens_with_style = [] | |
| for run in runs_in_paragraph: | |
| tokens = word_tokenize(run["text"]) | |
| for token in tokens: | |
| tokens_with_style.append(run.copy()) | |
| tokens_with_style[-1]["text"] = token | |
| token_index = 0 | |
| tokenized_sentences_with_style = [] | |
| for sentence in tokenized_sentences: | |
| sentence_with_style = [] | |
| for word in sentence: | |
| if word == tokens_with_style[token_index]["text"]: | |
| sentence_with_style.append(tokens_with_style[token_index]) | |
| token_index += 1 | |
| else: | |
| if word.startswith(tokens_with_style[token_index]["text"]): | |
| # this token might be split into several runs | |
| word_left = word | |
| while word_left: | |
| sentence_with_style.append(tokens_with_style[token_index]) | |
| word_left = word_left.removeprefix(tokens_with_style[token_index]["text"]) | |
| token_index += 1 | |
| else: | |
| raise "oops" | |
| tokenized_sentences_with_style.append(sentence_with_style) | |
| return tokenized_sentences_with_style | |
| def generate_alignments(original_runs_in_paragraph, translated_paragraph, aligner, temp_folder, detokenizer): | |
| # clean temp folder | |
| for f in os.listdir(temp_folder): | |
| os.remove(os.path.join(temp_folder, f)) | |
| temp_file_path = os.path.join(temp_folder, "tokenized_sentences.txt") | |
| # tokenize the original text by sentence and words while keeping the style | |
| original_tokenized_sentences_with_style = tokenize_paragraph_with_runs(original_runs_in_paragraph, detokenizer) | |
| # tokenize the translated text by sentence and word | |
| translated_tokenized_sentences = [word_tokenize(sentence) for sentence in sent_tokenize(translated_paragraph)] | |
| # write the file that fastalign will use | |
| with open(temp_file_path, "w") as out_file: | |
| for original, translated in zip(original_tokenized_sentences_with_style, translated_tokenized_sentences): | |
| out_file.write(f"{" ".join(item["text"] for item in original)} ||| {" ".join(translated)}\n") | |
| alignments = aligner.align(temp_file_path) | |
| # using the alignments generated by fastalign, we need to copy the style of the original token to the translated one | |
| translated_sentences_with_style = [] | |
| for sentence_idx, sentence_alignments in enumerate(alignments): | |
| # reverse the order of the alignments and build a dict with it | |
| sentence_alignments = {target: source for source, target in sentence_alignments} | |
| translated_sentence_with_style = [] | |
| for token_idx, translated_token in enumerate(translated_tokenized_sentences[sentence_idx]): | |
| # fastalign has found a token aligned with the translated one | |
| if token_idx in sentence_alignments.keys(): | |
| # get the aligned token | |
| original_idx = sentence_alignments[token_idx] | |
| new_entry = original_tokenized_sentences_with_style[sentence_idx][original_idx].copy() | |
| new_entry["text"] = translated_token | |
| translated_sentence_with_style.append(new_entry) | |
| else: | |
| # WARNING this is a test | |
| # since fastalign doesn't know from which word to reference this token, copy the style of the previous word | |
| new_entry = translated_sentence_with_style[-1].copy() | |
| new_entry["text"] = translated_token | |
| translated_sentence_with_style.append(new_entry) | |
| translated_sentences_with_style.append(translated_sentence_with_style) | |
| 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 = [] | |
| for key, group in groupby(values, key=lambda x: ( | |
| x['bold'], x['italic'], x['underline'], x['font_name'], x['font_size'], x['font_color'])): | |
| text = detokenizer.detokenize([item['text'] for item in group]) | |
| if groups and not text.startswith((",", ";", ":", ".", ")")): | |
| text = " " + text | |
| groups.append({"text": text, | |
| "bold": key[0], | |
| "italic": key[1], | |
| "underline": key[2], | |
| "font_name": key[3], | |
| "font_size": key[4], | |
| "font_color": key[5]}) | |
| return groups | |
| def preprocess_runs(runs_in_paragraph): | |
| new_runs = [] | |
| for run in runs_in_paragraph: | |
| # sometimes the parameters are False and sometimes they are None, set them all to False | |
| for key, value in run.items(): | |
| if value is None and not key.startswith("font"): | |
| run[key] = False | |
| if not new_runs: | |
| new_runs.append(run) | |
| else: | |
| # if the previous run has the same format as the current run, we merge the two runs together | |
| if (new_runs[-1]["bold"] == run["bold"] and new_runs[-1]["font_color"] == run["font_color"] and | |
| new_runs[-1]["font_color"] == run["font_color"] and new_runs[-1]["font_name"] == run["font_name"] | |
| and new_runs[-1]["font_size"] == run["font_size"] and new_runs[-1]["italic"] == run["italic"] | |
| and new_runs[-1]["underline"] == run["underline"]): | |
| new_runs[-1]["text"] += run["text"] | |
| else: | |
| new_runs.append(run) | |
| # we want to split runs that contain more than one sentence to avoid problems later when aligning styles | |
| sentences = sent_tokenize(new_runs[-1]["text"]) | |
| if len(sentences) > 1: | |
| new_runs[-1]["text"] = sentences[0] | |
| for sentence in sentences[1:]: | |
| new_run = new_runs[-1].copy() | |
| new_run["text"] = sentence | |
| new_runs.append(new_run) | |
| return new_runs | |
| if __name__ == "__main__": | |
| input_file = 'data/test3.docx' | |
| output_file = 'data/translated_output.docx' | |
| source_lang = 'ca' | |
| target_lang = 'en' | |
| config_folder = "fast_align_config" | |
| temp_folder = "tmp" | |
| aligner = Aligner(config_folder, source_lang, target_lang, temp_folder) | |
| os.makedirs(temp_folder, exist_ok=True) | |
| # load original file, extract the paragraphs with their runs (which include style and formatting) | |
| doc = Document(input_file) | |
| paragraphs_with_runs = extract_paragraphs_with_runs(doc) | |
| detokenizer = TreebankWordDetokenizer() | |
| # translate each paragraph | |
| translated_paragraphs = [] | |
| for paragraph in paragraphs_with_runs: | |
| paragraph_text = detokenizer.detokenize([run["text"] for run in paragraph]) | |
| translated_paragraphs.append(translate_paragraph(paragraph_text)) | |
| out_doc = Document() | |
| for original_runs_in_paragraph, translated_paragraph, original_paragraph in zip(paragraphs_with_runs, | |
| translated_paragraphs, | |
| doc.paragraphs): | |
| # in case there are empty paragraphs | |
| if len(original_runs_in_paragraph) == 1 and not original_runs_in_paragraph[0]["text"]: | |
| out_doc.add_paragraph(style=original_paragraph.style) | |
| original_runs_in_paragraph = preprocess_runs(original_runs_in_paragraph) | |
| paragraph_with_style = generate_alignments(original_runs_in_paragraph, translated_paragraph, aligner, | |
| temp_folder, detokenizer) | |
| para = out_doc.add_paragraph(style=original_paragraph.style) | |
| # flatten the paragraph, we don't need it to split into sentences anymore | |
| paragraph_with_style = [item for sublist in paragraph_with_style for item in sublist] | |
| # merge tokens into runs and detokenize | |
| paragraph_with_runs = group_by_style(paragraph_with_style, detokenizer) | |
| for item in paragraph_with_runs: | |
| run = para.add_run(item["text"]) | |
| # Preserve original run formatting | |
| run.bold = item['bold'] | |
| run.italic = item['italic'] | |
| run.underline = item['underline'] | |
| run.font.name = item['font_name'] | |
| run.font.size = item['font_size'] | |
| run.font.color.rgb = item['font_color'] | |
| out_doc.save(output_file) | |