Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sacremoses/truecase.py: 27%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# -*- coding: utf-8 -*-
3from __future__ import print_function
5import os
6import re
7from collections import defaultdict, Counter
8from functools import partial
9from itertools import chain
11from sacremoses.corpus import Perluniprops
12from sacremoses.util import parallelize_preprocess, grouper
15perluniprops = Perluniprops()
18class MosesTruecaser(object):
19 """
20 This is a Python port of the Moses Truecaser from
21 https://github.com/moses-smt/mosesdecoder/blob/master/scripts/recaser/train-truecaser.perl
22 https://github.com/moses-smt/mosesdecoder/blob/master/scripts/recaser/truecase.perl
23 """
25 #: Upper bound on distinct lowercased types read from a model file. A
26 #: model is often downloaded or shared, and loading amplifies its size
27 #: about 35x in memory, so this is a backstop against a hostile one -- not
28 #: a limit on legitimate vocabularies, which sit far below it.
29 MAX_MODEL_ENTRIES = 2_000_000
31 # Perl Unicode Properties character sets.
32 Lowercase_Letter = str("".join(perluniprops.chars("Lowercase_Letter")))
33 Uppercase_Letter = str("".join(perluniprops.chars("Uppercase_Letter")))
34 Titlecase_Letter = str("".join(perluniprops.chars("Uppercase_Letter")))
36 def __init__(self, load_from=None, is_asr=None, encoding="utf8"):
37 """
38 :param load_from:
39 :type load_from:
41 :param is_asr: A flag to indicate that model is for ASR. ASR input has
42 no case, make sure it is lowercase, and make sure known are cased
43 eg. 'i' to be uppercased even if i is known.
44 :type is_asr: bool
45 """
46 # Initialize the object.
47 super(MosesTruecaser, self).__init__()
48 # Initialize the language specific nonbreaking prefixes.
49 self.SKIP_LETTERS_REGEX = re.compile(
50 "[{}{}{}]".format(
51 self.Lowercase_Letter, self.Uppercase_Letter, self.Titlecase_Letter
52 )
53 )
55 self.XML_SPLIT_REGX = re.compile("(<.*(?<=>))(.*)((?=</)[^>]*>)")
57 self.SENT_END = {".", ":", "?", "!"}
58 self.DELAYED_SENT_START = {
59 "(",
60 "[",
61 '"',
62 "'",
63 "'",
64 """,
65 "[",
66 "]",
67 }
69 self.encoding = encoding
71 self.is_asr = is_asr
72 if load_from:
73 self.model = self._load_model(load_from)
75 def learn_truecase_weights(self, tokens, possibly_use_first_token=False):
76 """
77 This function checks through each tokens in a sentence and returns the
78 appropriate weight of each surface token form.
79 """
80 # Keep track of first tokens in the sentence(s) of the line.
81 is_first_word = True
82 truecase_weights = []
83 for i, token in enumerate(tokens):
84 # Skip XML tags.
85 if re.search(r"(<\S[^>]*>)", token):
86 continue
87 # Skip if sentence start symbols.
88 elif token in self.DELAYED_SENT_START:
89 continue
91 # Resets the `is_first_word` after seeing sent end symbols.
92 if not is_first_word and token in self.SENT_END:
93 is_first_word = True
94 continue
95 # Skips tokens with nothing to case.
96 if not self.SKIP_LETTERS_REGEX.search(token):
97 is_first_word = False
98 continue
100 # If it's not the first word,
101 # then set the current word weight to 1.
102 current_word_weight = 0
103 if not is_first_word:
104 current_word_weight = 1
105 # Otherwise check whether user wants to optionally
106 # use the first word.
107 elif possibly_use_first_token:
108 # Gated special handling of first word of sentence.
109 # Check if first characer of token is lowercase.
110 if token[0].islower():
111 current_word_weight = 1
112 elif i == 1:
113 current_word_weight = 0.1
115 is_first_word = False
117 if current_word_weight > 0:
118 truecase_weights.append((token.lower(), token, current_word_weight))
119 return truecase_weights
121 def _train(
122 self,
123 document_iterator,
124 save_to=None,
125 possibly_use_first_token=False,
126 processes=1,
127 progress_bar=False,
128 ):
129 """
130 :param document_iterator: The input document, each outer list is a sentence,
131 the inner list is the list of tokens for each sentence.
132 :type document_iterator: iter(list(str))
134 :param possibly_use_first_token: When True, on the basis that the first
135 word of a sentence is always capitalized; if this option is provided then:
136 a) if a sentence-initial token is *not* capitalized, then it is counted, and
137 b) if a capitalized sentence-initial token is the only token of the segment,
138 then it is counted, but with only 10% of the weight of a normal token.
139 :type possibly_use_first_token: bool
141 :returns: A dictionary of the best, known objects as values from `_casing_to_model()`
142 :rtype: {'best': dict, 'known': Counter}
143 """
144 casing = defaultdict(Counter)
145 train_truecaser = partial(
146 self.learn_truecase_weights,
147 possibly_use_first_token=possibly_use_first_token,
148 )
149 token_weights = chain(
150 *parallelize_preprocess(
151 train_truecaser, document_iterator, processes, progress_bar=progress_bar
152 )
153 )
154 # Collect the token_weights from every sentence.
155 for lowercase_token, surface_token, weight in token_weights:
156 casing[lowercase_token][surface_token] += weight
158 # Save to file if specified.
159 if save_to:
160 self._save_model_from_casing(casing, save_to)
161 return self._casing_to_model(casing)
163 def train(
164 self,
165 documents,
166 save_to=None,
167 possibly_use_first_token=False,
168 processes=1,
169 progress_bar=False,
170 ):
171 """
172 Default duck-type of _train(), accepts list(list(str)) as input documents.
173 """
174 self.model = None # Clear the model first.
175 self.model = self._train(
176 documents,
177 save_to,
178 possibly_use_first_token,
179 processes,
180 progress_bar=progress_bar,
181 )
182 return self.model
184 def train_from_file(
185 self,
186 filename,
187 save_to=None,
188 possibly_use_first_token=False,
189 processes=1,
190 progress_bar=False,
191 ):
192 """
193 Duck-type of _train(), accepts a filename to read as a `iter(list(str))`
194 object.
195 """
196 with open(filename, encoding=self.encoding) as fin:
197 # document_iterator = map(str.split, fin.readlines())
198 document_iterator = (
199 line.split() for line in fin.readlines()
200 ) # Lets try a generator comprehension for Python2...
201 self.model = None # Clear the model first.
202 self.model = self._train(
203 document_iterator,
204 save_to,
205 possibly_use_first_token,
206 processes,
207 progress_bar=progress_bar,
208 )
209 return self.model
211 def train_from_file_object(
212 self,
213 file_object,
214 save_to=None,
215 possibly_use_first_token=False,
216 processes=1,
217 progress_bar=False,
218 ):
219 """
220 Duck-type of _train(), accepts a file object to read as a `iter(list(str))`
221 object.
222 """
223 # document_iterator = map(str.split, file_object.readlines())
224 document_iterator = (
225 line.split() for line in file_object.readlines()
226 ) # Lets try a generator comprehension for Python2...
227 self.model = None # Clear the model first.
228 self.model = self._train(
229 document_iterator,
230 save_to,
231 possibly_use_first_token,
232 processes,
233 progress_bar=progress_bar,
234 )
235 return self.model
237 def truecase(self, text, return_str=False, use_known=False):
238 """
239 Truecase a single sentence / line of text.
241 :param text: A single string, i.e. sentence text.
242 :type text: str
244 :param use_known: Use the known case if a word is a known word but not the first word.
245 :type use_known: bool
246 """
247 check_model_message = str(
248 "\nUse Truecaser.train() to train a model.\n"
249 "Or use Truecaser('modefile') to load a model."
250 )
251 # Not an assert: `python -O` strips it, and the caller would then get
252 # an opaque AttributeError deep inside the loop instead of this message.
253 if not hasattr(self, "model"):
254 raise ValueError(check_model_message)
255 # Keep track of first tokens in the sentence(s) of the line.
256 is_first_word = True
257 truecased_tokens = []
258 tokens = self.split_xml(text)
259 # best_cases = best_cases if best_cases else self.model['best']
260 # known_cases = known_cases if known_cases else self.model['known']
262 for i, token in enumerate(tokens):
264 # Append XML tags and continue
265 if re.search(r"(<\S[^>]*>)", token):
266 truecased_tokens.append(token)
267 continue
269 # Note this shouldn't happen other if | are escaped as |
270 # To make the truecaser resilient,
271 # we'll just any token starting with pipes as they are.
272 if token == "|" or token.startswith("|"):
273 truecased_tokens.append(token)
274 continue
276 # Reads the word token and factors separatedly
277 token, other_factors = re.search(r"^([^\|]+)(.*)", token).groups()
279 # Lowercase the ASR tokens.
280 if self.is_asr:
281 token = token.lower()
283 # The actual case replacement happens here.
284 # "Most frequent" case of the word.
285 best_case = self.model["best"].get(token.lower(), None)
286 # If it's the start of sentence.
287 if is_first_word and best_case: # Truecase sentence start.
288 token = best_case
289 elif use_known and token in self.model["known"]: # Don't change known tokens.
290 pass
291 elif best_case: # Truecase otherwise unknown tokens? Heh? From https://github.com/moses-smt/mosesdecoder/blob/master/scripts/recaser/truecase.perl#L66
292 token = best_case
293 # Else, it's an unknown word, don't change the word.
294 # Concat the truecased `word` with the `other_factors`
295 token = token + other_factors
296 # Adds the truecased word.
297 truecased_tokens.append(token)
299 # Resets sentence start if this token is an ending punctuation.
300 if token in self.SENT_END:
301 is_first_word = True
302 elif token not in self.DELAYED_SENT_START:
303 is_first_word = False
305 # return ' '.join(tokens)
306 return " ".join(truecased_tokens) if return_str else truecased_tokens
308 def truecase_file(self, filename, return_str=True):
309 with open(filename, encoding=self.encoding) as fin:
310 for line in fin:
311 truecased_tokens = self.truecase(line.strip())
312 # Yield the truecased line.
313 yield " ".join(truecased_tokens) if return_str else truecased_tokens
315 @staticmethod
316 def split_xml(line):
317 """
318 Python port of split_xml function in Moses' truecaser:
319 https://github.com/moses-smt/mosesdecoder/blob/master/scripts/recaser/truecaser.perl
321 :param line: Input string, should be tokenized, separated by space.
322 :type line: str
323 """
324 line = line.strip()
325 tokens = []
326 while line:
327 # Assumes that xml tag is always separated by space.
328 has_xml = re.search(r"^\s*(<\S[^>]*>)(.*)$", line)
329 # non-XML test.
330 is_non_xml = re.search(r"^\s*([^\s<>]+)(.*)$", line)
331 # '<' or '>' occurs in word, but it's not an XML tag
332 xml_cognates = re.search(r"^\s*(\S+)(.*)$", line)
333 if has_xml:
334 potential_xml, line_next = has_xml.groups()
335 # exception for factor that is an XML tag
336 if (
337 re.search(r"^\S", line)
338 and len(tokens) > 0
339 and re.search(r"\|$", tokens[-1])
340 ):
341 tokens[-1] += potential_xml
342 # If it's a token with factors, join with the previous token.
343 is_factor = re.search(r"^(\|+)(.*)$", line_next)
344 if is_factor:
345 tokens[-1] += is_factor.group(1)
346 line_next = is_factor.group(2)
347 else:
348 tokens.append(
349 potential_xml + " "
350 ) # Token hack, unique to sacremoses.
351 line = line_next
353 elif is_non_xml:
354 tokens.append(is_non_xml.group(1)) # Token hack, unique to sacremoses.
355 line = is_non_xml.group(2)
356 elif xml_cognates:
357 tokens.append(
358 xml_cognates.group(1)
359 ) # Token hack, unique to sacremoses.
360 line = xml_cognates.group(2)
361 else:
362 raise Exception("ERROR: huh? {}".format(line))
363 tokens[-1] = tokens[-1].strip() # Token hack, unique to sacremoses.
364 return tokens
366 def _casing_to_model(self, casing):
367 """
369 :returns: A tuple of the (best, known) objects.
370 :rtype: tuple(dict, Counter)
371 """
372 best = {}
373 known = Counter()
375 for token_lower in casing:
376 tokens = casing[token_lower].most_common()
377 # Set the most frequent case as the "best" case.
378 best[token_lower] = tokens[0][0]
379 # If it's asr, throw away everything
380 if not self.is_asr:
381 for token, count in tokens[1:]:
382 # Note: This is rather odd that the counts are thrown away...
383 # from https://github.com/moses-smt/mosesdecoder/blob/master/scripts/recaser/truecase.perl#L34
384 known[token] += 1
385 model = {"best": best, "known": known, "casing": casing}
386 return model
388 def save_model(self, filename):
389 self._save_model_from_casing(self.model["casing"], filename)
391 def _save_model_from_casing(self, casing, filename):
392 """
393 Outputs the truecaser model file in the same output format as
394 https://github.com/moses-smt/mosesdecoder/blob/master/scripts/recaser/train-truecaser.perl
396 :param casing: The dictionary of tokens counter from `train()`.
397 :type casing: default(Counter)
398 """
399 # Write a fresh sibling file and rename it over the target, rather than
400 # truncating the target in place. Three things follow from that: a
401 # symlink at `filename` is replaced instead of being followed to
402 # whatever it points at (CWE-59), a crash mid-write cannot leave a
403 # half-written model behind, and the check-then-write race in the CLI
404 # ("if not os.path.isfile(modelfile): ... save_model(modelfile)")
405 # can no longer clobber someone else's file. O_EXCL|O_NOFOLLOW on the
406 # temporary name makes the create itself unspoofable; 0600 keeps the
407 # corpus vocabulary out of a world-readable file (CWE-732).
408 directory = os.path.dirname(os.path.abspath(filename))
409 tmpname = os.path.join(
410 directory, ".%s.%d.tmp" % (os.path.basename(filename), os.getpid())
411 )
412 flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
413 flags |= getattr(os, "O_NOFOLLOW", 0) # not defined on Windows
414 try:
415 handle = os.open(tmpname, flags, 0o600)
416 except FileExistsError:
417 os.unlink(tmpname)
418 handle = os.open(tmpname, flags, 0o600)
419 try:
420 with open(handle, "w", encoding=self.encoding) as fout:
421 self._write_casing(casing, fout)
422 os.replace(tmpname, filename)
423 except BaseException:
424 # Never leave the scratch file behind on an error path.
425 try:
426 os.unlink(tmpname)
427 except OSError:
428 pass
429 raise
431 def _write_casing(self, casing, fout):
432 for token in casing:
433 total_token_count = sum(casing[token].values())
434 tokens_counts = []
435 for i, (token, count) in enumerate(casing[token].most_common()):
436 if i == 0:
437 out_token = "{} ({}/{})".format(token, count, total_token_count)
438 else:
439 out_token = "{} ({})".format(token, count, total_token_count)
440 tokens_counts.append(out_token)
441 print(" ".join(tokens_counts), end="\n", file=fout)
443 def _load_model(self, filename):
444 """
445 Loads pre-trained truecasing file.
447 :returns: A dictionary of the best, known objects as values from `_casing_to_model()`
448 :rtype: {'best': dict, 'known': Counter}
449 """
450 casing = defaultdict(Counter)
451 with open(filename, encoding=self.encoding) as fin:
452 for lineno, line in enumerate(fin, 1):
453 line = line.strip().split()
454 for token, count in grouper(line, 2):
455 # `grouper` pads a short final pair with None, so a line
456 # with an odd number of fields used to surface as
457 # `AttributeError: 'NoneType' object has no attribute
458 # 'split'`, and a non-numeric count as a bare ValueError
459 # from int(). A downloaded/shared .truemodel is untrusted
460 # input, so report where it is malformed instead.
461 if count is None:
462 raise ValueError(
463 "malformed truecase model %r: line %d has an odd "
464 "number of fields (expected 'token (count/total)' "
465 "pairs)" % (filename, lineno)
466 )
467 try:
468 count = int(count.split("/")[0].strip("()"))
469 except ValueError:
470 raise ValueError(
471 "malformed truecase model %r: line %d has a "
472 "non-integer count %r" % (filename, lineno, count)
473 ) from None
474 casing[token.lower()][token] = count
475 if len(casing) > self.MAX_MODEL_ENTRIES:
476 # Loading costs roughly 35x the file size in RSS, so an
477 # oversized model is a memory-exhaustion vector (CWE-400).
478 # Raise MAX_MODEL_ENTRIES if you genuinely have a bigger
479 # vocabulary than this.
480 raise ValueError(
481 "truecase model %r exceeds MAX_MODEL_ENTRIES (%d) at "
482 "line %d" % (filename, self.MAX_MODEL_ENTRIES, lineno)
483 )
484 # Returns the best and known object from `_casing_to_model()`
485 return self._casing_to_model(casing)
488class MosesDetruecaser(object):
489 def __init__(self):
490 # Initialize the object.
491 super(MosesDetruecaser, self).__init__()
492 self.SENT_END = {".", ":", "?", "!"}
493 self.DELAYED_SENT_START = {
494 "(",
495 "[",
496 '"',
497 "'",
498 "'",
499 """,
500 "[",
501 "]",
502 }
504 # Some predefined tokens that will always be in lowercase.
505 self.ALWAYS_LOWER = {
506 "a",
507 "after",
508 "against",
509 "al-.+",
510 "and",
511 "any",
512 "as",
513 "at",
514 "be",
515 "because",
516 "between",
517 "by",
518 "during",
519 "el-.+",
520 "for",
521 "from",
522 "his",
523 "in",
524 "is",
525 "its",
526 "last",
527 "not",
528 "of",
529 "off",
530 "on",
531 "than",
532 "the",
533 "their",
534 "this",
535 "to",
536 "was",
537 "were",
538 "which",
539 "will",
540 "with",
541 }
543 def detruecase(self, text, is_headline=False, return_str=False):
544 """
545 Detruecase the translated files from a model that learnt from truecased
546 tokens.
548 :param text: A single string, i.e. sentence text.
549 :type text: str
550 """
551 # `cased_tokens` keep tracks of detruecased tokens.
552 cased_tokens = []
553 sentence_start = True
554 # Capitalize token if it's at the sentence start.
555 for token in text.split():
556 token = token[:1].upper() + token[1:] if sentence_start else token
557 cased_tokens.append(token)
558 if token in self.SENT_END:
559 sentence_start = True
560 elif not token in self.DELAYED_SENT_START:
561 sentence_start = False
562 # Check if it's a headline, if so then use title case.
563 if is_headline:
564 cased_tokens = [
565 token if token in self.ALWAYS_LOWER else token[:1].upper() + token[1:]
566 for token in cased_tokens
567 ]
569 return " ".join(cased_tokens) if return_str else cased_tokens
572__all__ = ["MosesTruecaser", "MosesDetruecaser"]