Add segmenter

app.py

@@ -1,13 +1,11 @@
-import gc
 import os
 from dotenv import load_dotenv
 import gradio as gr
 from AinaTheme import AinaGradioTheme
-import sentencepiece as spm
-import ctranslate2
 from huggingface_hub import snapshot_download
 import nltk
-from nltk import sent_tokenize
+
+from translate import translate_text
 
 nltk.download('punkt')
 
@@ -15,7 +13,7 @@ load_dotenv()
 
 MODELS_PATH = "./models"
 HF_CACHE_DIR = "./hf_cache"
-MAX_INPUT_CHARACTERS= int(os.environ.get("MAX_INPUT_CHARACTERS", default=500))
+MAX_INPUT_CHARACTERS= int(os.environ.get("MAX_INPUT_CHARACTERS", default=1000))
 
 def download_model(repo_id, revision="main"):
     return snapshot_download(repo_id=repo_id, revision=revision, local_dir=os.path.join(MODELS_PATH, repo_id), cache_dir=HF_CACHE_DIR)
@@ -30,14 +28,14 @@ model_dir_ca_fr = download_model("projecte-aina/mt-aina-ca-fr", revision="main")
 model_dir_fr_ca = download_model("projecte-aina/mt-aina-fr-ca", revision="main")
 
 model_dir_ca_de = download_model("projecte-aina/mt-aina-ca-de", revision="main")
-model_dir_de_ca = download_model("projecte-aina/mt-aina-de-ca", revision="main")
 
 model_dir_ca_it = download_model("projecte-aina/mt-aina-ca-it", revision="main")
-model_dir_it_ca = download_model("projecte-aina/mt-aina-it-ca", revision="main")
 
 model_dir_ca_pt = download_model("projecte-aina/mt-aina-ca-pt", revision="main")
 model_dir_pt_ca = download_model("projecte-aina/mt-aina-pt-ca", revision="main")
 
+model_dir_ca_zh = download_model("projecte-aina/mt-aina-ca-zh", revision="main")
+
 directions = {
     "Catalan": {
         "target": {
@@ -45,6 +43,7 @@ directions = {
             "English": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_en)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_en)}")},
             "French": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_fr)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_fr)}")},
             "German": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_de)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_de)}")},
+            "Chinese": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_zh)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_zh)}")},
             "Italian": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_it)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_it)}")},
             "Portuguese": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_pt)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_pt)}")}
 
@@ -65,16 +64,6 @@ directions = {
             "Catalan": {"model": (f"{os.path.join(MODELS_PATH, model_dir_fr_ca)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_fr_ca)}")},
         }
     },
-    "German": {
-        "target": {
-            "Catalan": {"model": (f"{os.path.join(MODELS_PATH, model_dir_de_ca)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_de_ca)}")},
-        }
-    },
-    "Italian": {
-        "target": {
-            "Catalan": {"model": (f"{os.path.join(MODELS_PATH, model_dir_it_ca)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_it_ca)}")},
-        }
-    },
     "Portuguese": {
         "target": {
             "Catalan": {"model": (f"{os.path.join(MODELS_PATH, model_dir_pt_ca)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_pt_ca)}")},
@@ -92,7 +81,7 @@ def get_target_languege_model(source_language, target_language):
     return directions.get(source_language, {}).get("target", {}).get(target_language, {}).get("model")
 
 
-def translate(source, lang_pair):
+def translate(text, source_language, lang_pair):
     """Use CTranslate model to translate a sentence
 
     Args:
@@ -102,18 +91,8 @@ def translate(source, lang_pair):
     Returns:
         Translation of the source text
     """
-    sp_model = spm.SentencePieceProcessor(lang_pair[0])
-    translator = ctranslate2.Translator(lang_pair[1])
-
-    source_sentences = sent_tokenize(source)
-    source_tokenized = sp_model.encode(source_sentences, out_type=str)
-    translations = translator.translate_batch(source_tokenized)
-    translations = [translation[0]["tokens"] for translation in translations]
-    translations_detokenized = sp_model.decode(translations)
-    translation = " ".join(translations_detokenized)
-    translation = translation.replace(' ⁇', ':')
-
-    gc.collect()
+  
+    translation = translate_text(text, source_language, lang_pair)
 
     return translation
 
@@ -125,7 +104,7 @@ def translate_input(input, source_language, target_language):
         return None
     
     target_language_model = get_target_languege_model(source_language, target_language)
-    translation = translate(input, target_language_model)
+    translation = translate(input, source_language, target_language_model)
 
     return translation
 

requirements.txt

@@ -3,4 +3,6 @@ gradio==4.8.0
 ctranslate2==3.23.0
 nltk==3.8.1
 sentencepiece==0.1.99
-python-dotenv==1.0.0
+python-dotenv==1.0.0
+pyonmttok==1.37.1
+lxml==4.9.3

srx_segmenter.py

@@ -0,0 +1,108 @@
+
+"""Segment text with SRX.
+"""
+__version__ = '0.0.2'
+
+import lxml.etree
+import regex
+from typing import (
+    List,
+    Set,
+    Tuple,
+    Dict,
+    Optional
+)
+
+
+class SrxSegmenter:
+    """Handle segmentation with SRX regex format.
+    """
+    def __init__(self, rule: Dict[str, List[Tuple[str, Optional[str]]]], source_text: str) -> None:
+        self.source_text = source_text
+        self.non_breaks = rule.get('non_breaks', [])
+        self.breaks = rule.get('breaks', [])
+
+    def _get_break_points(self, regexes: List[Tuple[str, str]]) -> Set[int]:
+        return set([
+            match.span(1)[1]
+            for before, after in regexes
+            for match in regex.finditer('({})({})'.format(before, after), self.source_text)
+        ])
+
+    def get_non_break_points(self) -> Set[int]:
+        """Return segment non break points
+        """
+        return self._get_break_points(self.non_breaks)
+
+    def get_break_points(self) -> Set[int]:
+        """Return segment break points
+        """
+        return self._get_break_points(self.breaks)
+
+    def extract(self) -> Tuple[List[str], List[str]]:
+        """Return segments and whitespaces.
+        """
+        non_break_points = self.get_non_break_points()
+        candidate_break_points = self.get_break_points()
+
+        break_point = sorted(candidate_break_points - non_break_points)
+        source_text = self.source_text
+
+        segments = []  # type: List[str]
+        whitespaces = []  # type: List[str]
+        previous_foot = ""
+        for start, end in zip([0] + break_point, break_point + [len(source_text)]):
+            segment_with_space = source_text[start:end]
+            candidate_segment = segment_with_space.strip()
+            if not candidate_segment:
+                previous_foot += segment_with_space
+                continue
+
+            head, segment, foot = segment_with_space.partition(candidate_segment)
+
+            segments.append(segment)
+            whitespaces.append('{}{}'.format(previous_foot, head))
+            previous_foot = foot
+        whitespaces.append(previous_foot)
+
+        return segments, whitespaces
+
+
+def parse(srx_filepath: str) -> Dict[str, Dict[str, List[Tuple[str, Optional[str]]]]]:
+    """Parse SRX file and return it.
+    :param srx_filepath: is soruce SRX file.
+    :return: dict
+    """
+    tree = lxml.etree.parse(srx_filepath)
+    namespaces = {
+        'ns': 'http://www.lisa.org/srx20'
+    }
+
+    rules = {}
+
+    for languagerule in tree.xpath('//ns:languagerule', namespaces=namespaces):
+        rule_name = languagerule.attrib.get('languagerulename')
+        if rule_name is None:
+            continue
+
+        current_rule = {
+            'breaks': [],
+            'non_breaks': [],
+        }
+
+        for rule in languagerule.xpath('ns:rule', namespaces=namespaces):
+            is_break = rule.attrib.get('break', 'yes') == 'yes'
+            rule_holder = current_rule['breaks'] if is_break else current_rule['non_breaks']
+
+            beforebreak = rule.find('ns:beforebreak', namespaces=namespaces)
+            beforebreak_text = '' if beforebreak.text is None else beforebreak.text
+
+            afterbreak = rule.find('ns:afterbreak', namespaces=namespaces)
+            afterbreak_text = '' if afterbreak.text is None else afterbreak.text
+
+            rule_holder.append((beforebreak_text, afterbreak_text))
+
+        rules[rule_name] = current_rule
+
+    #_add_softcatala_rules(rules)
+    return rules

texttokenizer.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+# -*- encoding: utf-8 -*-
+#
+# Copyright (c) 2020 Jordi Mas i Hernandez <jmas@softcatala.org>
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with this program; if not, write to the
+# Free Software Foundation, Inc., 59 Temple Place - Suite 330,
+# Boston, MA 02111-1307, USA.
+
+from __future__ import print_function
+from srx_segmenter import SrxSegmenter, parse
+import os
+
+
+def add_breakline_rule(rules,language):
+    rules[language]["breaks"].append(["\n", #Before
+                                      ""] # After
+                                     )
+    return rules
+
+
+'''
+    Splits text into sentences keeping spaces to allow later
+    to reconstruct the same text but with translatabled text changed 
+'''
+class TextTokenizer:
+    def __init__(self, language):
+        srx_filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'segment.srx')
+        self.rules = parse(srx_filepath)
+        self.language = language 
+        self.rules = add_breakline_rule(self.rules,language)
+       
+    def tokenize(self, sentence):
+        strings = []
+        translate = []
+
+        segmenter = SrxSegmenter(self.rules[self.language], sentence)
+        segments, whitespaces = segmenter.extract()
+
+        for i in range(len(segments)):
+            whitespace = whitespaces[i]
+            if len(whitespace) > 0:
+                strings.append(whitespace)
+                translate.append(False)
+
+            string = segments[i]
+            strings.append(string)
+            translate.append(True)
+
+        return strings, translate
+
+    def sentence_from_tokens(self, sentences, translate, translated):
+        num_sentences = len(sentences)
+        translation = ''
+        for i in range(0, num_sentences):
+            if translate[i] is True:
+                translation += translated[i]
+            else:
+                translation += sentences[i]
+
+        return translation.strip()

translate.py

@@ -0,0 +1,67 @@
+import gc
+import ctranslate2
+import pyonmttok
+from huggingface_hub import snapshot_download
+
+from texttokenizer import TextTokenizer
+import unicodedata
+
+
+def _normalize_input_string(result):
+    result = unicodedata.normalize('NFC', result)
+    return result
+
+def _translate_batch(input_batch, spm, model, max_sentence_batch=10):
+
+    batch_input_tokenized = []
+    batch_input_markers = []
+
+    #preserve_markup = PreserveMarkup()
+
+    num_sentences = len(input_batch)
+    for pos in range(0, num_sentences):
+        tokenized = spm.tokenize(input_batch[pos])[0]
+        batch_input_tokenized.append(tokenized)
+
+    batch_output = []
+    for offset in range(0,len(batch_input_tokenized), max_sentence_batch):
+      partial_result = model.translate_batch(batch_input_tokenized[offset:offset+max_sentence_batch], return_scores=False, replace_unknowns=True)
+      for pos in range(0,len(partial_result)):
+        tokenized = partial_result[pos][0]['tokens']
+        translated = spm.detokenize(tokenized)
+        batch_output.append(translated)
+
+    return batch_output
+
+
+def translate_text(sample_text, source_language, lang_pair, max_sentence_batch=20):
+     
+	spm = pyonmttok.Tokenizer(mode="none",sp_model_path=lang_pair[0])
+	translator = ctranslate2.Translator(lang_pair[1], device="cpu")
+	tokenizer=TextTokenizer(source_language)
+
+	text = _normalize_input_string(sample_text)
+	sentences, translate = tokenizer.tokenize(text)
+	num_sentences = len(sentences)
+	sentences_batch = []
+	indexes = []
+	results = ["" for x in range(num_sentences)]
+	for i in range(num_sentences):
+		if translate[i] is False:
+			continue
+
+		sentences_batch.append(sentences[i])
+		indexes.append(i)
+
+	translated_batch = _translate_batch(sentences_batch, spm, translator, max_sentence_batch)
+	for pos in range(0, len(translated_batch)):
+		i = indexes[pos]
+		results[i] = translated_batch[pos]
+
+	#Rebuild split sentences
+	translated = tokenizer.sentence_from_tokens(sentences, translate, results)
+     
+	gc.collect()
+     
+	return translated
+