init commit

AUTHORS

@@ -0,0 +1,12 @@
+# This is the list of Ithaca authors for copyright purposes.
+#
+# This does not necessarily list everyone who has contributed code, since in
+# some cases, their employer may be the copyright holder.  To see the full list
+# of contributors, see the revision history in source control.
+DeepMind Technologies Limited
+Google LLC
+Thea Sommerschield
+Jonathan Prag
+Marita Chatzipanagiotou
+John Pavlopoulos
+Ion Androutsopoulos

CONTRIBUTING.md

@@ -0,0 +1,29 @@
+# How to Contribute
+
+We'd love to accept your patches and contributions to this project. There are
+just a few small guidelines you need to follow.
+
+## Contributor License Agreement
+
+Contributions to this project must be accompanied by a Contributor License
+Agreement (CLA). You (or your employer) retain the copyright to your
+contribution; this simply gives us permission to use and redistribute your
+contributions as part of the project. Head over to
+<https://cla.developers.google.com/> to see your current agreements on file or
+to sign a new one.
+
+You generally only need to submit a CLA once, so if you've already submitted one
+(even if it was for a different project), you probably don't need to do it
+again.
+
+## Code reviews
+
+All submissions, including submissions by project members, require review. We
+use GitHub pull requests for this purpose. Consult
+[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
+information on using pull requests.
+
+## Community Guidelines
+
+This project follows
+[Google's Open Source Community Guidelines](https://opensource.google/conduct/).

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

example_input.txt

@@ -0,0 +1 @@
+ἔδοξεν τῆι βουλῆι καὶ τῶι δήμωι λυσίστρατος εἶπε- ἐπειδὴ διοφάνης ἀνὴρ ἀγαθὸς ὢν διατελεῖ περὶ δηλίους δεδόχθαι τῶι ----- διοφάνην καλλι-------- --ηναῖον πρόξενον εἶναι δ--------- αὐτὸγ καὶ ἐκγόνους κ-- εἶναι αὐτοῖς ἀτέλειαν ἐν δήλωι πάντων καὶ γῆς καὶ οἰκίας ἔγκτησιν καὶ πρόσοδον πρὸς τὴμ βουλὴγ καὶ τὸν δῆμον πρώτοις μετὰ τὰ ἱερὰ καὶ τὰ ἄλλα ὅσα καὶ τοῖς ἄλλοις προξένοις καὶ εὐεργέταις τοῦ ἱεροῦ δέδοται παρὰ ---ίων ἀναγράψαι δὲ τόδε ?????????α τὴν βουλὴν εἰς -----------ριον τοὺς -ὲ -----------------------.

inference_example.py

@@ -0,0 +1,129 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Example for running inference. See also colab."""
+
+import functools
+import pickle
+
+from absl import app
+from absl import flags
+from ithaca.eval import inference
+from ithaca.models.model import Model
+from ithaca.util.alphabet import GreekAlphabet
+import jax
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    'input', '', 'Text to directly pass to the model. Only one of --input and '
+    '--input_file can be specified.')
+flags.DEFINE_string(
+    'input_file', '', 'File containing text to pass to the model. Only one of '
+    '--input and --input_file can be specified.')
+flags.DEFINE_string('checkpoint_path', 'checkpoint.pkl',
+                    'Path to model checkpoint pickle.')
+flags.DEFINE_string('attribute_json', '', 'Path to save attribution JSON to.')
+flags.DEFINE_string('restore_json', '', 'Path to save restoration JSON to.')
+
+
+def load_checkpoint(path):
+  """Loads a checkpoint pickle.
+
+  Args:
+    path: path to checkpoint pickle
+
+  Returns:
+    a model config dictionary (arguments to the model's constructor), a dict of
+    dicts containing region mapping information, a GreekAlphabet instance with
+    indices and words populated from the checkpoint, a dict of Jax arrays
+    `params`, and a `forward` function.
+  """
+
+  # Pickled checkpoint dict containing params and various config:
+  with open(path, 'rb') as f:
+    checkpoint = pickle.load(f)
+
+  # We reconstruct the model using the same arguments as during training, which
+  # are saved as a dict in the "model_config" key, and construct a `forward`
+  # function of the form required by attribute() and restore().
+  params = jax.device_put(checkpoint['params'])
+  model = Model(**checkpoint['model_config'])
+  forward = functools.partial(model.apply, params)
+
+  # Contains the mapping between region IDs and names:
+  region_map = checkpoint['region_map']
+
+  # Use vocabulary mapping from the checkpoint, the rest of the values in the
+  # class are fixed and constant e.g. the padding symbol
+  alphabet = GreekAlphabet()
+  alphabet.idx2word = checkpoint['alphabet']['idx2word']
+  alphabet.word2idx = checkpoint['alphabet']['word2idx']
+
+  return checkpoint['model_config'], region_map, alphabet, params, forward
+
+
+def main(argv):
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  if FLAGS.input and not FLAGS.input_file:
+    input_text = FLAGS.input
+  elif not FLAGS.input and FLAGS.input_file:
+    with open(FLAGS.input_file, 'r', encoding='utf8') as f:
+      input_text = f.read()
+  else:
+    raise app.UsageError('Specify exactly one of --input and --input_file.')
+
+  if not 50 <= len(input_text) <= 750:
+    raise app.UsageError(
+        f'Text should be between 50 and 750 chars long, but the input was '
+        f'{len(input_text)} characters')
+
+  # Load the checkpoint pickle and extract from it the pieces needed for calling
+  # the attribute() and restore() functions:
+  (model_config, region_map, alphabet, params,
+   forward) = load_checkpoint(FLAGS.checkpoint_path)
+  vocab_char_size = model_config['vocab_char_size']
+  vocab_word_size = model_config['vocab_word_size']
+
+  attribution = inference.attribute(
+      input_text,
+      forward=forward,
+      params=params,
+      alphabet=alphabet,
+      region_map=region_map,
+      vocab_char_size=vocab_char_size,
+      vocab_word_size=vocab_word_size)
+  if FLAGS.attribute_json:
+    with open(FLAGS.attribute_json, 'w') as f:
+      f.write(attribution.json(indent=2))
+  else:
+    print('Attribution:', attribution.json())
+
+  restoration = inference.restore(
+      input_text,
+      forward=forward,
+      params=params,
+      alphabet=alphabet,
+      vocab_char_size=vocab_char_size,
+      vocab_word_size=vocab_word_size)
+  if FLAGS.restore_json:
+    with open(FLAGS.restore_json, 'w') as f:
+      f.write(restoration.json(indent=2))
+  else:
+    print('Restoration:', restoration.json())
+
+
+if __name__ == '__main__':
+  app.run(main)

ithaca/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

ithaca/eval/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

ithaca/eval/inference.py

@@ -0,0 +1,268 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for performing inference using Jax, including decoding.
+
+The module is separated into two main entrypoints: attribute() and restore().
+
+Both take a function called `forward`, a Jax function mapping from model inputs
+(excluding parameters) to the model output tuple. Generated using
+e.g. `functools.partial(exp.forward.apply, exp._params)`.
+"""
+
+import json
+import math
+import re
+from typing import List, NamedTuple, Tuple
+
+import ithaca.util.eval as eval_util
+import ithaca.util.text as util_text
+
+import jax
+import numpy as np
+
+
+class LocationPrediction(NamedTuple):
+  """One location prediction and its associated probability."""
+
+  location_id: int
+  score: float
+
+  def build_json(self):
+    return {
+        'location_id': self.location_id,
+        'score': self.score,
+    }
+
+
+class AttributionResults(NamedTuple):
+  """Immediate model output attribution predictions and related information."""
+
+  input_text: str
+
+  # List of pairs of location ID and probability
+  locations: List[LocationPrediction]
+
+  # Probabilities over year range [-800, -790, -780, ..., 790, 800]
+  year_scores: List[float]  # length 160
+
+  # Per-character saliency maps:
+  date_saliency: List[float]
+  location_saliency: List[float]  # originally called subregion
+
+  def build_json(self):
+    return {
+        'input_text': self.input_text,
+        'locations': [l.build_json() for l in self.locations],
+        'year_scores': self.year_scores,
+        'date_saliency': self.date_saliency,
+        'location_saliency': self.location_saliency
+    }
+
+  def json(self, **kwargs):
+    return json.dumps(self.build_json(), **kwargs)
+
+
+class Restoration(NamedTuple):
+  """One restored candidate string from the beam search."""
+  text: str
+  score: float
+
+  def build_json(self):
+    return {'text': self.text, 'score': self.score}
+
+
+class RestorationCharSaliency(NamedTuple):
+  """Saliency entry for one predicted character of a prediction."""
+  text: str
+  restored_idx: int  # which predicted character the saliency map corresponds to
+  saliency: List[float]
+
+  def build_json(self):
+    return {
+        'text': self.text,
+        'restored_idx': self.restored_idx,
+        'saliency': self.saliency
+    }
+
+
+class RestorationResults(NamedTuple):
+  """Contains all text-related restoration predictions."""
+
+  input_text: str
+  top_prediction: str
+  restored: List[int]  # char indices that were missing (-)
+
+  # List of top N results from beam search:
+  predictions: List[Restoration]
+
+  # Saliency maps for each successive character of the best (greedy) prediction
+  prediction_saliency: List[RestorationCharSaliency]
+
+  def build_json(self):
+    return {
+        'input_text':
+            self.input_text,
+        'top_prediction':
+            self.top_prediction,
+        'restored':
+            self.restored,
+        'predictions': [r.build_json() for r in self.predictions],
+        'prediction_saliency': [
+            m.build_json() for m in self.prediction_saliency
+        ],
+    }
+
+  def json(self, **kwargs):
+    return json.dumps(self.build_json(), **kwargs)
+
+
+# These constants are fixed for all recent versions of the model.
+MIN_TEXT_LEN = 50
+TEXT_LEN = 768  # fixed model sequence length
+DATE_MIN = -800
+DATE_MAX = 800
+DATE_INTERVAL = 10
+RESTORATION_BEAM_WIDTH = 20
+RESTORATION_TEMPERATURE = 1.
+SEED = 1
+ALPHABET_MISSING_RESTORE = '?'  # missing characters to restore
+
+
+def _prepare_text(
+    text, alphabet
+) -> Tuple[str, str, str, np.ndarray, np.ndarray, List[int], np.ndarray,
+           List[int]]:
+  """Adds start of sequence symbol, and padding.
+
+  Also strips accents if present, trims whitespace, and generates arrays ready
+  for input into the model.
+
+  Args:
+    text: Raw text input string, no padding or start of sequence symbol.
+    alphabet: GreekAlphabet object containing index/character mappings.
+
+  Returns:
+    Tuple of cleaned text (str), padded text (str), char indices (array of batch
+    size 1), word indices (array of batch size 1), text length (list of size 1)
+  """
+  text = re.sub(r'\s+', ' ', text.strip())
+  text = util_text.strip_accents(text)
+
+  if len(text) < MIN_TEXT_LEN:
+    raise ValueError('Input text too short.')
+
+  if len(text) >= TEXT_LEN - 1:
+    raise ValueError('Input text too long.')
+
+  text_sos = alphabet.sos + text
+  text_len = [len(text_sos)]  # includes SOS, but not padding
+
+  text_padded = text_sos + alphabet.pad * max(0, TEXT_LEN - len(text_sos))
+
+  restore_mask_idx = [
+      i for i, c in enumerate(text_padded) if c == ALPHABET_MISSING_RESTORE
+  ]
+  text_padded = text_padded.replace(ALPHABET_MISSING_RESTORE, alphabet.missing)
+
+  text_char = util_text.text_to_idx(text_padded, alphabet).reshape(1, -1)
+  text_word = util_text.text_to_word_idx(text_padded, alphabet).reshape(1, -1)
+  padding = np.where(text_char > 0, 1, 0)
+
+  return (text, text_sos, text_padded, text_char, text_word, text_len, padding,
+          restore_mask_idx)
+
+
+def attribute(text, forward, params, alphabet, vocab_char_size, vocab_word_size,
+              region_map) -> AttributionResults:
+  """Computes predicted date and geographical region."""
+
+  (text, _, _, text_char, text_word, text_len, padding,
+   _) = _prepare_text(text, alphabet)
+
+  rng = jax.random.PRNGKey(SEED)
+  date_logits, subregion_logits, _, _ = forward(
+      text_char=text_char,
+      text_word=text_word,
+      rngs={'dropout': rng},
+      is_training=False)
+
+  # Generate subregion predictions:
+  subregion_logits = np.array(subregion_logits)
+  subregion_pred_probs = eval_util.softmax(subregion_logits[0]).tolist()
+  location_predictions = [
+      LocationPrediction(location_id=id, score=prob)
+      for prob, id in zip(subregion_pred_probs, region_map['sub']['ids'])
+  ]
+  location_predictions.sort(key=lambda loc: loc.score, reverse=True)
+
+  # Generate date predictions:
+  date_pred_probs = eval_util.softmax(date_logits[0])
+
+  # Gradients for saliency maps
+  date_saliency, subregion_saliency = eval_util.compute_attribution_saliency_maps(
+      text_char, text_word, text_len, padding, forward, params, rng, alphabet,
+      vocab_char_size, vocab_word_size)
+
+  # Skip start of sequence symbol (first char) for text and saliency maps:
+  return AttributionResults(
+      input_text=text,
+      locations=location_predictions,
+      year_scores=date_pred_probs.tolist(),
+      date_saliency=date_saliency.tolist()[1:],
+      location_saliency=subregion_saliency.tolist()[1:])
+
+
+def restore(text, forward, params, alphabet, vocab_char_size,
+            vocab_word_size) -> RestorationResults:
+  """Performs search to compute text restoration. Slower, runs synchronously."""
+
+  if ALPHABET_MISSING_RESTORE not in text:
+    raise ValueError('At least one character must be missing.')
+
+  text, _, text_padded, _, _, text_len, _, restore_mask_idx = _prepare_text(
+      text, alphabet)
+
+  beam_result = eval_util.beam_search_batch_2d(
+      forward,
+      alphabet,
+      text_padded,
+      restore_mask_idx,
+      beam_width=RESTORATION_BEAM_WIDTH,
+      temperature=RESTORATION_TEMPERATURE,
+      rng=jax.random.PRNGKey(SEED))
+
+  # For visualization purposes, we strip out the SOS and padding, and adjust
+  # restored_indices accordingly
+  predictions = [
+      Restoration(
+          text=beam_entry.text_pred[1:].rstrip(alphabet.pad),
+          score=math.exp(beam_entry.pred_logprob)) for beam_entry in beam_result
+  ]
+  restored_indices = [i - 1 for i in restore_mask_idx]
+
+  # Sequence of saliency maps for a greedy prediction:
+  saliency_steps = eval_util.sequential_restoration_saliency(
+      text_padded, text_len, forward, params, alphabet, restore_mask_idx,
+      vocab_char_size, vocab_word_size)
+
+  return RestorationResults(
+      input_text=text,
+      top_prediction=predictions[0].text,
+      restored=restored_indices,
+      predictions=predictions,
+      prediction_saliency=[
+          RestorationCharSaliency(step.text, int(step.pred_char_pos),
+                                  step.saliency_map.tolist())
+          for step in saliency_steps
+      ])

ithaca/models/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

ithaca/models/bigbird.py

@@ -0,0 +1,110 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformer using BigBird (https://arxiv.org/abs/2007.14062).
+
+This implementation is from the Long Range Arena:
+https://github.com/google-research/long-range-arena/tree/main/lra_benchmarks/models/bigbird
+"""
+
+from typing import Any, Optional
+
+from . import bigbird_attention
+from . import common_layers
+
+from flax import linen as nn
+import jax.numpy as jnp
+
+_DEFAULT_BLOCK_SIZE = 64
+_DEFAULT_NUM_RAND_BLOCKS = 3
+
+
+class BigBirdBlock(nn.Module):
+  """BigBird layer (https://arxiv.org/abs/2007.14062).
+
+  Attributes:
+    qkv_dim: dimension of the query/key/value
+    mlp_dim: dimension of the mlp on top of attention block
+    num_heads: number of heads
+    dtype: the dtype of the computation (default: float32).
+    causal_mask: bool, mask future or not
+    dropout_rate: dropout rate
+    attention_dropout_rate: dropout rate for attention weights
+    deterministic: bool, deterministic or not (to apply dropout)
+    activation_fn: Activation function ("relu", "gelu")
+    block_size: Size of attention blocks.
+    num_rand_blocks: Number of random blocks.
+    connectivity_seed: Optional seed for random block sparse attention.
+  """
+
+  qkv_dim: Any
+  mlp_dim: Any
+  num_heads: Any
+  dtype: Any = jnp.float32
+  causal_mask: bool = False
+  dropout_rate: float = 0.1
+  attention_dropout_rate: float = 0.1
+  deterministic: bool = False
+  activation_fn: str = 'relu'
+  block_size: int = _DEFAULT_BLOCK_SIZE
+  num_rand_blocks: int = _DEFAULT_NUM_RAND_BLOCKS
+  connectivity_seed: Optional[int] = None
+
+  @nn.compact
+  def __call__(self, inputs, inputs_segmentation=None, padding_mask=None):
+    """Applies BigBirdBlock module.
+
+    Args:
+      inputs: input data
+      inputs_segmentation: input segmentation info for packed examples.
+      padding_mask: bool, mask padding tokens, [b, l, 1]
+
+    Returns:
+      output after transformer block.
+
+    """
+
+    # Attention block.
+    assert inputs.ndim == 3
+    x = common_layers.LayerNorm(dtype=self.dtype)(inputs)
+    x = bigbird_attention.BigBirdSelfAttention(
+        num_heads=self.num_heads,
+        dtype=self.dtype,
+        qkv_features=self.qkv_dim,
+        kernel_init=nn.initializers.xavier_uniform(),
+        bias_init=nn.initializers.normal(stddev=1e-6),
+        use_bias=False,
+        broadcast_dropout=False,
+        dropout_rate=self.attention_dropout_rate,
+        deterministic=self.deterministic,
+        block_size=self.block_size,
+        num_rand_blocks=self.num_rand_blocks,
+        connectivity_seed=self.connectivity_seed)(
+            x,
+            segmentation=inputs_segmentation,
+            padding_mask=padding_mask,
+        )
+    x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=self.deterministic)
+    x = x + inputs
+
+    # MLP block.
+    y = common_layers.LayerNorm(dtype=self.dtype)(x)
+    y = common_layers.MlpBlock(
+        mlp_dim=self.mlp_dim,
+        dtype=self.dtype,
+        dropout_rate=self.dropout_rate,
+        deterministic=self.deterministic,
+        activation_fn=self.activation_fn)(
+            y)
+
+    return x + y

ithaca/models/bigbird_attention.py

@@ -0,0 +1,602 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Big Bird attention mechanism.
+
+See https://arxiv.org/abs/2007.14062.
+
+This implementation is from the Long Range Arena:
+https://github.com/google-research/long-range-arena/tree/main/lra_benchmarks/models/bigbird
+"""
+
+# pylint: disable=attribute-defined-outside-init,g-bare-generic
+import functools
+from typing import Any, Callable, Optional
+from absl import logging
+from flax import linen as nn
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+def get_block_rand_mask(m, n, wm, wn, r, last_idx=-1):
+  """This function creates the m by n mask for random block sparse mask.
+
+  Args:
+    m: input size
+    n: output size
+    wm: block input size
+    wn: block output size
+    r: number of random block per row
+    last_idx: if -1 then r blocks are chosen throughout the n space, if
+      possitive then r blocks are chooses at random upto last_ids
+
+  Returns:
+    blocked mask of size m//wm -2 by r
+  """
+  if (m // wm) != (n // wn):
+    logging.info('Error the number of blocks needs to be same')
+  rand_attn = np.zeros((m // wm - 2, r), dtype=jnp.int64)
+  a = np.array(range(1, n // wn - 1))
+  last = (m // wn) - 1
+  if last_idx > (2 * wn):
+    last = (last_idx // wn) - 1
+  for i in range(1, m // wm - 1):
+    start = i - 2
+    end = i
+    if i == 1:
+      rand_attn[i - 1, :] = np.random.permutation(a[2:last])[:r]
+    elif i == 2:
+      rand_attn[i - 1, :] = np.random.permutation(a[3:last])[:r]
+    elif i == m // wm - 3:
+      rand_attn[i - 1, :] = np.random.permutation(a[:last - 4])[:r]
+    elif i == m // wm - 2:
+      rand_attn[i - 1, :] = np.random.permutation(a[:last - 3])[:r]
+    else:
+      if start > last:
+        start = last
+        rand_attn[i - 1, :] = np.random.permutation(a[:start])[:r]
+      elif (end + 1) == last:
+        rand_attn[i - 1, :] = np.random.permutation(a[:start])[:r]
+      else:
+        rand_attn[i - 1, :] = np.random.permutation(
+            np.concatenate((a[:start], a[end + 1:last])))[:r]
+  return rand_attn
+
+
+def create_band_mask_from_inputs(from_blocked_mask, to_blocked_mask):
+  """Create 3D attention mask from a 2D tensor mask.
+
+  Args:
+    from_blocked_mask: 2D Tensor of shape [batch_size,
+      from_seq_length//from_block_size, from_block_size].
+    to_blocked_mask: int32 Tensor of shape [batch_size,
+      to_seq_length//to_block_size, to_block_size].
+
+  Returns:
+    float Tensor of shape [batch_size, 1, from_seq_length//from_block_size-4,
+                           from_block_size,  3*to_block_size].
+  """
+  exp_blocked_to_pad = jnp.concatenate([
+      to_blocked_mask[:, 1:-3], to_blocked_mask[:, 2:-2], to_blocked_mask[:,
+                                                                          3:-1]
+  ], 2)
+  band_pad = jnp.einsum('BLQ,BLK->BLQK', from_blocked_mask[:, 2:-2],
+                        exp_blocked_to_pad)
+  band_pad = jnp.expand_dims(band_pad, 1)
+  return band_pad
+
+
+def create_rand_mask_from_inputs(from_blocked_mask, to_blocked_mask, rand_attn):
+  """Create 3D attention mask from a 2D tensor mask.
+
+  Args:
+    from_blocked_mask: 2D Tensor of shape [batch_size,
+      from_seq_length//from_block_size, from_block_size].
+    to_blocked_mask: int32 Tensor of shape [batch_size,
+      to_seq_length//to_block_size, to_block_size].
+    rand_attn: [batch_size, num_attention_heads,
+      from_seq_length//from_block_size-2, rsize]
+
+  Returns:
+    float Tensor of shape [batch_size, num_attention_heads,
+    from_seq_length//from_block_size-2,
+                           from_block_size,  3*to_block_size].
+  """
+
+  # batch_size, num_attention_heads, num_windows, _ = get_shape_list(
+  #     rand_attn, expected_rank=4)
+  batch_size, num_attention_heads, num_windows, _ = rand_attn.shape
+  rand_pad = jnp.reshape(
+      # Equivalent to tf.gather(to_blocked_mask, rand_attn, batch_dims=1)
+      gather_1(to_blocked_mask, rand_attn),
+      [batch_size, num_attention_heads, num_windows, -1])
+  rand_pad = jnp.einsum('BLQ,BHLK->BHLQK', from_blocked_mask[:, 1:-1], rand_pad)
+  return rand_pad
+
+
+@jax.vmap
+def gather_1(params, indices):
+  return jnp.take(params, indices, axis=0)
+
+
+@jax.vmap
+def gather_2(params, indices):
+  return gather_1(params, indices)
+
+
+def band_start_block_rand_multi_attention_pad(query_matrix, key_matrix,
+                                              value_matrix, rand_attn, band_pad,
+                                              rand_pad, seq_m_pad, seq_n_pad, b,
+                                              h, m, wm, n, wn, r, d):
+  """Applies sparse block band rand attention in hopefully efficient way.
+
+  Args:
+    query_matrix: b, h, n, d
+    key_matrix: b, h, n, d
+    value_matrix: b, h, n, d
+    rand_attn: b, h, m//wm-2, r
+    band_pad: b, 1, m//wm-4, wm, 3*wn
+    rand_pad: b, h, m//wm-2, wm, r*wn
+    seq_m_pad: b, 1, m, 1
+    seq_n_pad: b, 1, 1, n
+    b: batch size
+    h: number of head
+    m: from_length
+    wm: from window size
+    n: to length
+    wn: to window size
+    r: number of rand blocks
+    d: hidden dimension
+
+  Returns:
+    context layer. b, m, h, -1
+    attention weights. [b, h, m//wm-4, wm, (5+r)*wn]
+  """
+  blocked_query_matrix = jnp.reshape(query_matrix, (b, h, m // wm, wm, -1))
+  blocked_key_matrix = jnp.reshape(key_matrix, (b, h, n // wn, wn, -1))
+  blocked_value_matrix = jnp.reshape(value_matrix, (b, h, n // wn, wn, -1))
+  # tf.gather(blocked_key_matrix, rand_attn, batch_dims=2, name='gather_key'),
+  gathered_key = jnp.reshape(
+      gather_2(blocked_key_matrix, rand_attn),
+      (b, h, m // wm - 2, r * wn, -1))  # [b, h, n//wn-2, r, wn, -1]
+  # tf.gather(
+  #   blocked_value_matrix, rand_attn, batch_dims=2, name='gather_value')
+  gathered_value = jnp.reshape(
+      gather_2(blocked_value_matrix, rand_attn),
+      (b, h, m // wm - 2, r * wn, -1))  # [b, h, n//wn-2, r, wn, -1]
+
+  first_product = jnp.einsum(
+      'BHQD,BHKD->BHQK', blocked_query_matrix[:, :, 0],
+      key_matrix)  # [b, h, wm, -1] x [b, h, n, -1] ==> [b, h, wm, n]
+  first_product = first_product / jnp.sqrt(d)
+  first_product += (1.0 - seq_n_pad) * -10000.0
+  first_attn_weights = jax.nn.softmax(first_product)  # [b, h, wm, n]
+  first_context_layer = jnp.einsum(
+      'BHQK,BHKD->BHQD', first_attn_weights,
+      value_matrix)  # [b, h, wm, n] x [b, h, n, -1] ==> [b, h, wm, -1]
+  first_context_layer = jnp.expand_dims(first_context_layer, 2)
+
+  second_key_mat = jnp.concatenate([
+      blocked_key_matrix[:, :, 0], blocked_key_matrix[:, :, 1],
+      blocked_key_matrix[:, :, 2], blocked_key_matrix[:, :,
+                                                      -1], gathered_key[:, :, 0]
+  ], 2)  # [b, h, (4+r)*wn, -1]
+  second_value_mat = jnp.concatenate([
+      blocked_value_matrix[:, :, 0], blocked_value_matrix[:, :, 1],
+      blocked_value_matrix[:, :, 2], blocked_value_matrix[:, :, -1],
+      gathered_value[:, :, 0]
+  ], 2)  # [b, h, (4+r)*wn, -1]
+  second_product = jnp.einsum(
+      'BHQD,BHKD->BHQK', blocked_query_matrix[:, :, 1], second_key_mat
+  )  # [b, h, wm, -1] x [b, h, (4+r)*wn, -1] ==> [b, h, wm, (4+r)*wn]
+  second_seq_pad = jnp.concatenate([
+      seq_n_pad[:, :, :, :3 * wn], seq_n_pad[:, :, :, -wn:],
+      jnp.ones([b, 1, 1, r * wn], dtype=jnp.float32)
+  ], 3)
+  second_rand_pad = jnp.concatenate(
+      [jnp.ones([b, h, wm, 4 * wn], dtype=jnp.float32), rand_pad[:, :, 0]], 3)
+  second_product = second_product / jnp.sqrt(d)
+  second_product += (1.0 -
+                     jnp.minimum(second_seq_pad, second_rand_pad)) * -10000.0
+  second_attn_weights = jax.nn.softmax(second_product)  # [b , h, wm, (4+r)*wn]
+  second_context_layer = jnp.einsum(
+      'BHQK,BHKD->BHQD', second_attn_weights, second_value_mat
+  )  # [b, h, wm, (4+r)*wn] x [b, h, (4+r)*wn, -1] ==> [b, h, wm, -1]
+  second_context_layer = jnp.expand_dims(second_context_layer, 2)
+
+  exp_blocked_key_matrix = jnp.concatenate([
+      blocked_key_matrix[:, :, 1:-3], blocked_key_matrix[:, :, 2:-2],
+      blocked_key_matrix[:, :, 3:-1]
+  ], 3)  # [b, h, m//wm-4, 3*wn, -1]
+  exp_blocked_value_matrix = jnp.concatenate([
+      blocked_value_matrix[:, :, 1:-3], blocked_value_matrix[:, :, 2:-2],
+      blocked_value_matrix[:, :, 3:-1]
+  ], 3)  # [b, h, m//wm-4, 3*wn, -1]
+  middle_query_matrix = blocked_query_matrix[:, :, 2:-2]
+  inner_band_product = jnp.einsum(
+      'BHLQD,BHLKD->BHLQK', middle_query_matrix, exp_blocked_key_matrix
+  )  # [b, h, m//wm-4, wm, -1] x [b, h, m//wm-4, 3*wn, -1]
+  #     ==> [b, h, m//wm-4, wm, 3*wn]
+  inner_band_product = inner_band_product / jnp.sqrt(d)
+  rand_band_product = jnp.einsum(
+      'BHLQD,BHLKD->BHLQK', middle_query_matrix,
+      gathered_key[:, :,
+                   1:-1])  # [b, h, m//wm-4, wm, -1] x [b, h, m//wm-4, r*wn, -1]
+  #     ==> [b, h, m//wm-4, wm, r*wn]
+  rand_band_product = rand_band_product / jnp.sqrt(d)
+  first_band_product = jnp.einsum(
+      'BHLQD,BHKD->BHLQK', middle_query_matrix, blocked_key_matrix[:, :, 0]
+  )  # [b, h, m//wm-4, wm, -1] x [b, h, wn, -1] ==> [b, h, m//wm-4, wm, wn]
+  first_band_product = first_band_product / jnp.sqrt(d)
+  last_band_product = jnp.einsum(
+      'BHLQD,BHKD->BHLQK', middle_query_matrix, blocked_key_matrix[:, :, -1]
+  )  # [b, h, m//wm-4, wm, -1] x [b, h, wn, -1] ==> [b, h, m//wm-4, wm, wn]
+  last_band_product = last_band_product / jnp.sqrt(d)
+  inner_band_product += (1.0 - band_pad) * -10000.0
+  first_band_product += (1.0 -
+                         jnp.expand_dims(seq_n_pad[:, :, :, :wn], 3)) * -10000.0
+  last_band_product += (1.0 -
+                        jnp.expand_dims(seq_n_pad[:, :, :, -wn:], 3)) * -10000.0
+  rand_band_product += (1.0 - rand_pad[:, :, 1:-1]) * -10000.0
+  band_product = jnp.concatenate([
+      first_band_product, inner_band_product, rand_band_product,
+      last_band_product
+  ], -1)  # [b, h, m//wm-4, wm, (5+r)*wn]
+  attn_weights = jax.nn.softmax(band_product)  # [b, h, m//wm-4, wm, (5+r)*wn]
+  context_layer = jnp.einsum(
+      'BHLQK,BHLKD->BHLQD', attn_weights[:, :, :, :,
+                                         wn:4 * wn], exp_blocked_value_matrix
+  )  # [b, h, m//wm-4, wm, 3*wn] x [b, h, m//wm-4, 3*wn, -1]
+  #     ==> [b, h, m//wm-4, wm, -1]
+  context_layer += jnp.einsum(
+      'BHLQK,BHLKD->BHLQD', attn_weights[:, :, :, :,
+                                         4 * wn:-wn], gathered_value[:, :, 1:-1]
+  )  # [b, h, m//wm-4, wm, r*wn] x [b, h, m//wm-4, r*wn, -1]
+  #     ==> [b, h, m//wm-4, wm, -1]
+  context_layer += jnp.einsum(
+      'BHLQK,BHKD->BHLQD', attn_weights[:, :, :, :, :wn],
+      blocked_value_matrix[:, :, 0]
+  )  # [b, h, m//wm-4, wm, wn] x [b, h, wn, -1] ==> [b, h, m//wm-4, wm, -1]
+  context_layer += jnp.einsum(
+      'BHLQK,BHKD->BHLQD', attn_weights[:, :, :, :,
+                                        -wn:], blocked_value_matrix[:, :, -1]
+  )  # [b, h, m//wm-4, wm, wn] x [b, h, wn, -1] ==> [b, h, m//wm-4, wm, -1]
+
+  second_last_key_mat = jnp.concatenate([
+      blocked_key_matrix[:, :, 0], blocked_key_matrix[:, :, -3],
+      blocked_key_matrix[:, :, -2], blocked_key_matrix[:, :, -1],
+      gathered_key[:, :, -1]
+  ], 2)  # [b, h, (4+r)*wn, -1]
+  second_last_value_mat = jnp.concatenate([
+      blocked_value_matrix[:, :, 0], blocked_value_matrix[:, :, -3],
+      blocked_value_matrix[:, :, -2], blocked_value_matrix[:, :, -1],
+      gathered_value[:, :, -1]
+  ], 2)  # [b, h, (4+r)*wn, -1]
+  second_last_product = jnp.einsum(
+      'BHQD,BHKD->BHQK', blocked_query_matrix[:, :, -2], second_last_key_mat
+  )  # [b, h, wm, -1] x [b, h, (4+r)*wn, -1] ==> [b, h, wm, (4+r)*wn]
+  second_last_seq_pad = jnp.concatenate([
+      seq_n_pad[:, :, :, :wn], seq_n_pad[:, :, :, -3 * wn:],
+      jnp.ones([b, 1, 1, r * wn], dtype=jnp.float32)
+  ], 3)
+  second_last_rand_pad = jnp.concatenate(
+      [jnp.ones([b, h, wm, 4 * wn], dtype=jnp.float32), rand_pad[:, :, -1]], 3)
+  second_last_product = second_last_product / jnp.sqrt(d)
+  second_last_product += (
+      1.0 - jnp.minimum(second_last_seq_pad, second_last_rand_pad)) * -10000.0
+  second_last_attn_weights = jax.nn.softmax(
+      second_last_product)  # [b, h, wm, (4+r)*wn]
+  second_last_context_layer = jnp.einsum(
+      'BHQK,BHKD->BHQD', second_last_attn_weights, second_last_value_mat
+  )  # [b, h, wm, (4+r)*wn] x [b, h, (4+r)*wn, -1] ==> [b, h, wm, -1]
+  second_last_context_layer = jnp.expand_dims(second_last_context_layer, 2)
+
+  last_product = jnp.einsum(
+      'BHQD,BHKD->BHQK', blocked_query_matrix[:, :, -1],
+      key_matrix)  # [b, h, wm, -1] x [b, h, n, -1] ==> [b, h, wm, n]
+  last_product = last_product / jnp.sqrt(d)
+  last_product += (1.0 - seq_n_pad) * -10000.0
+  last_attn_weights = jax.nn.softmax(last_product)  # [b, h, wm, n]
+  last_context_layer = jnp.einsum(
+      'BHQK,BHKD->BHQD', last_attn_weights,
+      value_matrix)  # [b, h, wm, n] x [b, h, n, -1] ==> [b, h, wm, -1]
+  last_context_layer = jnp.expand_dims(last_context_layer, 2)
+
+  context_layer = jnp.concatenate([
+      first_context_layer, second_context_layer, context_layer,
+      second_last_context_layer, last_context_layer
+  ], 2)
+  context_layer = jnp.reshape(context_layer, (b, h, m, -1)) * seq_m_pad
+  context_layer = jnp.transpose(context_layer, (0, 2, 1, 3))
+  return context_layer, attn_weights
+
+
+def sparse_dot_product_attention(queries,
+                                 keys,
+                                 values,
+                                 connectivity_seed,
+                                 input_mask=None,
+                                 block_size=64,
+                                 num_rand_blocks=3):
+  """Implements sparse dot product attention given query, key, and value.
+
+  This is the core function for applying attention based on
+  https://arxiv.org/abs/1706.03762. It calculates the attention weights given
+  query and key and combines the values using the attention weights. This
+  function supports multi-dimensional inputs.
+
+
+  Args:
+    queries: queries for calculating attention with shape of `[batch_size,
+      length, num_heads, mem_channels]`.
+    keys: keys for calculating attention with shape of `[batch_size, length,
+      num_heads, mem_channels]`.
+    values: values to be used in attention with shape of `[batch_size, length,
+      num_heads, value_channels]`.
+    connectivity_seed: Integer seed for generating connectivity graph.
+    input_mask: Optional mask for keys/values with shape `[batch_size, length]`
+      and the same dtype.
+    block_size: Size for local attention around diagonal of attention.
+    num_rand_blocks: int. Number of random chunks per row.
+
+  Returns:
+    Output of shape `[bs, length, num_heads, value_channels]`.
+  """
+  (batch_size, to_seq_length, num_attention_heads, hidden_size) = keys.shape
+  from_seq_length = queries.shape[1]
+  seq_length = max(to_seq_length, from_seq_length)
+  queries = jnp.pad(queries,
+                    ((0, 0), (0, seq_length - from_seq_length), (0, 0), (0, 0)))
+  keys = jnp.pad(keys,
+                 ((0, 0), (0, seq_length - to_seq_length), (0, 0), (0, 0)))
+  values = jnp.pad(values,
+                   ((0, 0), (0, seq_length - to_seq_length), (0, 0), (0, 0)))
+
+  if input_mask is None:
+    input_mask = jnp.ones((batch_size, seq_length), dtype=keys.dtype)
+  else:
+    input_mask = jnp.pad(
+        input_mask,
+        tuple((0, seq_length - size) if i == 1 else (0, 0)
+              for i, size in enumerate(input_mask.shape)))
+
+  np.random.seed(connectivity_seed)
+  # pylint: disable=g-complex-comprehension
+  rand_attn = [
+      get_block_rand_mask(
+          seq_length,
+          seq_length,
+          block_size,
+          block_size,
+          num_rand_blocks,
+          last_idx=min(seq_length, 1024)) for _ in range(num_attention_heads)
+  ]
+  # pylint: enable=g-complex-comprehension
+  rand_attn = jnp.stack(rand_attn, axis=0)
+  rand_attn = jnp.expand_dims(rand_attn, 0)
+  rand_attn = jnp.repeat(rand_attn, batch_size, 0)
+
+  # reshape and cast for blocking
+  blocked_input_mask = jnp.reshape(
+      input_mask, (batch_size, seq_length // block_size, block_size))
+  input_mask = jnp.reshape(input_mask, (batch_size, 1, seq_length, 1))
+  output_mask = jnp.reshape(input_mask, (batch_size, 1, 1, seq_length))
+
+  # create band padding
+  band_pad = create_band_mask_from_inputs(blocked_input_mask,
+                                          blocked_input_mask)
+  rand_pad = create_rand_mask_from_inputs(blocked_input_mask,
+                                          blocked_input_mask, rand_attn)
+
+  queries = jnp.transpose(queries, (0, 2, 1, 3))
+  keys = jnp.transpose(keys, (0, 2, 1, 3))
+  values = jnp.transpose(values, (0, 2, 1, 3))
+
+  # sparse mask
+  context_layer, _ = band_start_block_rand_multi_attention_pad(
+      queries, keys, values, rand_attn, band_pad, rand_pad, input_mask,
+      output_mask, batch_size, num_attention_heads, seq_length, block_size,
+      seq_length, block_size, num_rand_blocks, hidden_size)
+
+  return context_layer[:, :from_seq_length, ...]
+
+
+class BigBirdAttention(nn.Module):
+  """Multi-head dot-product attention.
+
+  Attributes:
+    num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
+      should be divisible by the number of heads.
+    block_size: Size for local attention around diagonal of attention.
+    num_rand_blocks: int. Number of random chunks per row.
+    dtype: the dtype of the computation (default: float32)
+    qkv_features: dimension of the key, query, and value.
+    out_features: dimension of the last projection
+    broadcast_dropout: bool: use a broadcasted dropout along batch dims.
+    dropout_rate: dropout rate
+    deterministic: bool, deterministic or not (to apply dropout)
+    precision: numerical precision of the computation see `jax.lax.Precision`
+      for details.
+    kernel_init: initializer for the kernel of the Dense layers.
+    bias_init: initializer for the bias of the Dense layers.
+    use_bias: bool: whether pointwise QKVO dense transforms use bias.
+    connectivity_seed: Seed for random block sparse attention.
+  """
+
+  num_heads: int
+  block_size: int = 64
+  num_rand_blocks: int = 3
+  dtype: Any = jnp.float32
+  qkv_features: Optional[int] = None
+  out_features: Optional[int] = None
+  broadcast_dropout: bool = True
+  dropout_rate: float = 0.
+  deterministic: bool = False
+  precision: Any = None
+  kernel_init: Callable = nn.linear.default_kernel_init
+  bias_init: Callable = nn.initializers.zeros
+  use_bias: bool = True
+  connectivity_seed: Optional[int] = None
+
+  @nn.compact
+  def __call__(self,
+               inputs_q,
+               inputs_kv,
+               padding_mask=None,
+               segmentation=None,
+               dropout_rng=None):
+    """Applies multi-head dot product attention on the input data.
+
+    Projects the inputs into multi-headed query, key, and value vectors,
+    applies dot-product attention and project the results to an output vector.
+
+    This can be used for encoder-decoder attention by specifying both `inputs_q`
+    and `inputs_kv` orfor self-attention by only specifying `inputs_q` and
+    setting `inputs_kv` to None.
+
+    Args:
+      inputs_q: input queries of shape `[bs, length, features]`.
+      inputs_kv: key/values of shape `[bs, length, features]` or None for
+        self-attention, inn which case key/values will be derived from inputs_q.
+      padding_mask: boolean specifying query tokens that are pad token. [b, l,
+        1]
+      segmentation: segment indices for packed inputs_q data.
+      dropout_rng: JAX PRNGKey: to be used for dropout
+
+    Returns:
+      output of shape `[bs, length, features]`.
+    """
+
+    orig_seqlen = inputs_q.shape[-2]
+    extra_len = self.block_size - (orig_seqlen % self.block_size)
+    pad_width = np.array([[0, 0], [0, extra_len], [0, 0]])
+    mask_pad = np.array([[0, 0], [0, extra_len], [0, 0]])
+    padding_mask = jnp.pad(padding_mask, mask_pad, constant_values=-1e9)
+
+    inputs_q = jnp.pad(inputs_q, pad_width)
+    if inputs_kv is not None:
+      inputs_kv = jnp.pad(inputs_kv, pad_width)
+
+    if inputs_kv is None:
+      inputs_kv = inputs_q
+
+    features = self.out_features or inputs_q.shape[-1]
+    qkv_features = self.qkv_features or inputs_q.shape[-1]
+
+    assert qkv_features % self.num_heads == 0, (
+        'Memory dimension must be divisible by number of heads.')
+    head_dim = qkv_features // self.num_heads
+
+    dense = functools.partial(
+        nn.DenseGeneral,
+        axis=-1,
+        features=(self.num_heads, head_dim),
+        kernel_init=self.kernel_init,
+        bias_init=self.bias_init,
+        use_bias=self.use_bias,
+        precision=self.precision)
+    # project inputs_q to multi-headed q/k/v
+    # dimensions are then [bs, dims..., n_heads, n_features_per_head]
+    query, key, value = (dense(dtype=self.dtype, name='query')(inputs_q),
+                         dense(dtype=self.dtype, name='key')(inputs_kv),
+                         dense(dtype=self.dtype, name='value')(inputs_kv))
+
+    if self.connectivity_seed is None:
+      path = self._get_construction_frame().path
+      connectivity_seed = hash(path) % 2**32
+    else:
+      connectivity_seed = self.connectivity_seed
+    # apply attention
+    input_mask = None
+    if padding_mask is not None:
+      input_mask = padding_mask.astype(key.dtype)
+    x = sparse_dot_product_attention(
+        query,
+        key,
+        value,
+        connectivity_seed=connectivity_seed,
+        input_mask=input_mask,
+        block_size=self.block_size,
+        num_rand_blocks=self.num_rand_blocks)
+
+    # back to the original inputs dimensions
+    out = nn.DenseGeneral(
+        features=features,
+        axis=(-2, -1),
+        kernel_init=self.kernel_init,
+        bias_init=self.bias_init,
+        use_bias=self.use_bias,
+        dtype=self.dtype,
+        precision=self.precision,
+        name='out')(
+            x)
+
+    out = out[:, :orig_seqlen, :]
+
+    return out
+
+
+class BigBirdSelfAttention(BigBirdAttention):
+  """Multi-head dot-product self-attention.
+
+  Attributes:
+    num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])
+      should be divisible by the number of heads.
+    block_size: Size for local attention around diagonal of attention.
+    num_rand_blocks: int. Number of random chunks per row.
+    dtype: the dtype of the computation (default: float32)
+    qkv_features: dimension of the key, query, and value.
+    out_features: dimension of the last projection
+    broadcast_dropout: bool: use a broadcasted dropout along batch dims.
+    dropout_rate: dropout rate
+    deterministic: bool, deterministic or not (to apply dropout)
+    precision: numerical precision of the computation see `jax.lax.Precision`
+      for details.
+    kernel_init: initializer for the kernel of the Dense layers.
+    bias_init: initializer for the bias of the Dense layers.
+    use_bias: bool: whether pointwise QKVO dense transforms use bias.
+    connectivity_seed: Seed for random block sparse attention.
+  """
+
+  @nn.compact
+  def __call__(self,
+               inputs_q,
+               padding_mask=None,
+               segmentation=None,
+               dropout_rng=None):
+    """Applies multi-head dot product attention on the input data.
+
+    Projects the inputs into multi-headed query, key, and value vectors,
+    applies dot-product attention and project the results to an output vector.
+
+    This can be used for encoder-decoder attention by specifying both `inputs_q`
+    and `inputs_kv` orfor self-attention by only specifying `inputs_q` and
+    setting `inputs_kv` to None.
+
+    Args:
+      inputs_q: input queries of shape `[bs, length, features]`.
+      padding_mask: boolean specifying query tokens that are pad token.
+      segmentation: segment indices for packed inputs_q data.
+      dropout_rng: JAX PRNGKey: to be used for dropout
+
+    Returns:
+      output of shape `[bs, length, features]`.
+    """
+    return super().__call__(
+        inputs_q=inputs_q,
+        inputs_kv=None,
+        padding_mask=padding_mask,
+        segmentation=segmentation,
+        dropout_rng=dropout_rng,
+    )

ithaca/models/common_layers.py

@@ -0,0 +1,318 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Common layers used in models.
+
+This implementation is from the Long Range Arena:
+https://github.com/google-research/long-range-arena/tree/main/lra_benchmarks/models/bigbird
+"""
+
+# pylint: disable=attribute-defined-outside-init,g-bare-generic
+from typing import Any, Callable, Iterable, Optional
+
+from flax import linen as nn
+from jax import lax
+from jax.nn import initializers
+import jax.numpy as jnp
+import numpy as np
+
+PRNGKey = Any
+Array = Any
+Shape = Iterable[int]
+Dtype = Any  # this could be a real type?
+
+ACTIVATION_FN_DICT = {
+    'relu': nn.relu,
+    'gelu': nn.gelu,
+}
+
+
+def grid_restack(all_vecs):
+  """Grid restack for meta-performer.
+
+  Given multiple sequences (lists) of batch x len x dim,
+  reshape this such that all positions are side by side.
+
+  for example (for illustrative purposes):
+
+  inputs: [1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]
+  outputs: [1, 5, 9, 2, 6, 10, 3, 7, 11, 4, 8, 12]
+
+  Args:
+    all_vecs: list of sequences of batch x len x dim
+
+  Returns:
+    Array of batch x (length x num_items) x dim.
+  """
+  cat_output = []
+  for pos in range(all_vecs[0].shape[1]):
+    pos_vecs = [x[:, None, pos, :] for x in all_vecs]
+    cat_output += pos_vecs
+  x2 = jnp.concatenate(cat_output, 1)
+  return x2
+
+
+def shift_right(x):
+  """Shift the input to the right by padding on axis 1."""
+  pad_widths = [(0, 0)] * len(x.shape)
+  pad_widths[1] = (1, 0)  # Padding on axis=1
+  padded = jnp.pad(
+      x, pad_widths, mode='constant', constant_values=x.dtype.type(0))
+  return padded[:, :-1]
+
+
+class Embed(nn.Module):
+  """Embedding Module.
+
+  A parameterized function from integers [0, n) to d-dimensional vectors.
+
+  Attributes:
+    mode: either 'input' or 'output' -> to share input/output embedding
+    emb_init: embedding initializer
+  """
+
+  mode: str = 'input'
+  emb_init: Callable = nn.initializers.normal(stddev=1.0)
+
+  @nn.compact
+  def __call__(self, inputs, num_embeddings, features):
+    """Applies Embed module.
+
+    Args:
+      inputs: input data
+      num_embeddings: number of embedding
+      features: size of the embedding dimension
+
+    Returns:
+      output which is embedded input data
+    """
+    embedding = self.param('embedding', self.emb_init,
+                           (num_embeddings, features))
+    if self.mode == 'input':
+      if inputs.dtype not in [jnp.int32, jnp.int64, jnp.uint32, jnp.uint64]:
+        raise ValueError('Input type must be an integer or unsigned integer.')
+      return jnp.take(embedding, inputs, axis=0)
+    if self.mode == 'output':
+      return jnp.einsum('bld,vd->blv', inputs, embedding)
+
+
+def sinusoidal_init(max_len=2048, replicate_tf=False):
+  """1D Sinusoidal Position Embedding Initializer.
+
+  Args:
+      max_len: maximum possible length for the input
+      replicate_tf: replicate TF periodic encoding exactly
+
+  Returns:
+      output: init function returning `(1, max_len, d_feature)`
+  """
+
+  def init(key, shape, dtype=np.float32):
+    """Sinusoidal init."""
+    del key, dtype
+    d_feature = shape[-1]
+    pe = np.zeros((max_len, d_feature), dtype=np.float32)
+    position = np.arange(0, max_len)[:, np.newaxis]
+    if replicate_tf:
+      half_d_feature = d_feature // 2
+      div_term = np.exp(
+          np.arange(half_d_feature) * -(np.log(10000.0) / (half_d_feature - 1)))
+      pe[:, :half_d_feature] = np.sin(position * div_term)
+      pe[:, half_d_feature:] = np.cos(position * div_term)
+    else:
+      div_term = np.exp(
+          np.arange(0, d_feature, 2) * -(np.log(10000.0) / d_feature))
+      pe[:, 0::2] = np.sin(position * div_term)
+      pe[:, 1::2] = np.cos(position * div_term)
+    pe = pe[np.newaxis, :, :]  # [1, max_len, d_feature]
+    return jnp.array(pe)
+
+  return init
+
+
+class AddPositionEmbs(nn.Module):
+  """Adds (optionally learned) positional embeddings to the inputs.
+
+  Attributes:
+    posemb_init: positional embedding initializer, if None, then use a fixed
+      (non-learned) sinusoidal embedding table.
+    max_len: maximum possible length for the input.
+    replicate_original: replicate original periodic encoding exactly
+  """
+
+  posemb_init: Optional[Callable] = None
+  posemb_dim: Optional[int] = None
+  max_len: int = 512
+  combine_type: str = 'concat'
+  replicate_tf: bool = False
+
+  @nn.compact
+  def __call__(self, inputs, inputs_positions=None, cache=None):
+    """Applies AddPositionEmbs module.
+
+    By default this layer uses a fixed sinusoidal embedding table. If a
+    learned position embedding is desired, pass an initializer to
+    posemb_init.
+
+    Args:
+      inputs: input data.
+      inputs_positions: input position indices for packed sequences.
+      cache: flax attention cache for fast decoding.
+
+    Returns:
+      output: `(bs, timesteps, in_dim)`
+    """
+    # inputs.shape is (batch_size, seq_len, emb_dim)
+    assert inputs.ndim == 3, ('Number of dimensions should be 3,'
+                              ' but it is: %d' % inputs.ndim)
+    batch_size = inputs.shape[0]
+    length = inputs.shape[1]
+    if self.posemb_dim is None or self.combine_type == 'add':
+      self.posemb_dim = inputs.shape[-1]
+    pos_emb_shape = (1, self.max_len, self.posemb_dim)
+    if self.posemb_init is None:
+      # Use a fixed (non-learned) sinusoidal position embedding.
+      pos_embedding = sinusoidal_init(
+          max_len=self.max_len,
+          replicate_tf=self.replicate_tf,
+      )(None, pos_emb_shape, None)
+    else:
+      pos_embedding = self.param('pos_embedding', self.posemb_init,
+                                 pos_emb_shape)
+    pe = pos_embedding[:, :length, :]
+    # We abuse the same attention Cache mechanism to run positional embeddings
+    # in fast predict mode. We could use state variables instead, but this
+    # simplifies invocation with a single top-level cache context manager.
+    # We only use the cache's position index for tracking decoding position.
+    if cache:
+      if self.is_initializing():
+        cache.store(np.array((4, 1, 1), dtype=np.int32))
+      else:
+        cache_entry = cache.retrieve(None)
+        i = cache_entry.i
+        cache.store(cache_entry.replace(i=cache_entry.i + 1))
+        _, _, df = pos_embedding.shape
+        pe = lax.dynamic_slice(pos_embedding, jnp.array((0, i, 0)),
+                               jnp.array((1, 1, df)))
+    if inputs_positions is None:
+      # normal unpacked case:
+      if self.combine_type == 'add':
+        return inputs + pe
+      elif self.combine_type == 'concat':
+        pe_broadcast = np.repeat(pe, batch_size, axis=0)
+        return lax.concatenate([inputs, pe_broadcast], 2)
+      else:
+        raise ValueError('Wrong type value.')
+    else:
+      # for packed data we need to use known position indices:
+      return inputs + jnp.take(pe[0], inputs_positions, axis=0)
+
+
+class MlpBlock(nn.Module):
+  """Transformer MLP block."""
+
+  mlp_dim: int
+  dtype: Any = jnp.float32
+  out_dim: Optional[int] = None
+  out_dropout: bool = True
+  dropout_rate: float = 0.1
+  deterministic: bool = False
+  kernel_init: Callable = nn.initializers.xavier_uniform()
+  bias_init: Callable = nn.initializers.normal(stddev=1e-6)
+  activation_fn: str = 'gelu'
+
+  @nn.compact
+  def __call__(self, inputs):
+    """Applies Transformer MlpBlock module."""
+    actual_out_dim = inputs.shape[-1] if self.out_dim is None else self.out_dim
+    x = nn.Dense(
+        self.mlp_dim,
+        dtype=self.dtype,
+        kernel_init=self.kernel_init,
+        bias_init=self.bias_init)(
+            inputs)
+    x = ACTIVATION_FN_DICT[self.activation_fn](x)
+    x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=self.deterministic)
+    output = nn.Dense(
+        actual_out_dim,
+        dtype=self.dtype,
+        kernel_init=self.kernel_init,
+        bias_init=self.bias_init)(
+            x)
+    if self.out_dropout:
+      output = nn.Dropout(rate=self.dropout_rate)(
+          output, deterministic=self.deterministic)
+    return output
+
+
+def classifier_head(encoded, num_classes, mlp_dim, pooling_mode='MEAN'):
+  """Classifier head.
+
+  We put this here just so that all models consistently call the same function.
+
+  Args:
+    encoded: tensor inputs are shape of [bs, len, dim].
+    num_classes: int, number of classes
+    mlp_dim: int, dim of intermediate MLP.
+    pooling_mode: str, string dictating pooling op {MEAN}
+
+  Returns:
+    tensor of shape [bs, num_classes]
+
+  """
+  if pooling_mode == 'MEAN':
+    encoded = jnp.mean(encoded, axis=1)
+  elif pooling_mode == 'SUM':
+    encoded = jnp.sum(encoded, axis=1)
+  elif pooling_mode == 'FLATTEN':
+    encoded = encoded.reshape((encoded.shape[0], -1))
+  elif pooling_mode == 'CLS':
+    encoded = encoded[:, 0]
+  else:
+    raise NotImplementedError('Pooling not supported yet.')
+  encoded = nn.Dense(mlp_dim, name='mlp')(encoded)
+  encoded = nn.relu(encoded)
+  encoded = nn.Dense(num_classes, name='logits')(encoded)
+  return encoded
+
+
+class LayerNorm(nn.Module):
+  """Layer Norm to replicate tf.contrib."""
+  epsilon: Optional[float] = None
+  dtype: Any = jnp.float32
+  use_bias: bool = True
+  use_scale: bool = True
+  bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.zeros
+  scale_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.ones
+
+  @nn.compact
+  def __call__(self, x):
+    if self.epsilon is None:
+      epsilon = 1e-12 if self.dtype != jnp.float16 else 1e-3
+    else:
+      epsilon = self.epsilon
+    x = jnp.asarray(x, jnp.float32)
+    features = x.shape[-1]
+    mean = jnp.mean(x, axis=-1, keepdims=True)
+    mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
+    var = mean2 - lax.square(mean)
+    mul = lax.rsqrt(var + epsilon)
+    if self.use_scale:
+      mul = mul * jnp.asarray(
+          self.param('scale', self.scale_init, (features,)), self.dtype)
+    y = x * mul
+    if self.use_bias:
+      y = y + jnp.asarray(
+          self.param('bias', self.bias_init, (features,)), self.dtype)
+    y -= mean * mul
+    return jnp.asarray(y, self.dtype)

ithaca/models/model.py

@@ -0,0 +1,243 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Ithaca model."""
+
+from . import bigbird
+from . import common_layers
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+
+
+class Model(nn.Module):
+  """Transformer Model for sequence tagging."""
+  vocab_char_size: int = 164
+  vocab_word_size: int = 100004
+  output_subregions: int = 85
+  output_date: int = 160
+  output_date_dist: bool = True
+  output_return_emb: bool = False
+  use_output_mlp: bool = True
+  num_heads: int = 8
+  num_layers: int = 6
+  word_char_emb_dim: int = 192
+  emb_dim: int = 512
+  qkv_dim: int = 512
+  mlp_dim: int = 2048
+  max_len: int = 1024
+  causal_mask: bool = False
+  feature_combine_type: str = 'concat'
+  posemb_combine_type: str = 'add'
+  region_date_pooling: str = 'first'
+  learn_pos_emb: bool = True
+  use_bfloat16: bool = False
+  dropout_rate: float = 0.1
+  attention_dropout_rate: float = 0.1
+  activation_fn: str = 'gelu'
+  model_type: str = 'bigbird'
+
+  def setup(self):
+    self.text_char_emb = nn.Embed(
+        num_embeddings=self.vocab_char_size,
+        features=self.word_char_emb_dim,
+        embedding_init=nn.initializers.normal(stddev=1.0),
+        name='char_embeddings')
+    self.text_word_emb = nn.Embed(
+        num_embeddings=self.vocab_word_size,
+        features=self.word_char_emb_dim,
+        embedding_init=nn.initializers.normal(stddev=1.0),
+        name='word_embeddings')
+
+  @nn.compact
+  def __call__(self,
+               text_char=None,
+               text_word=None,
+               text_char_onehot=None,
+               text_word_onehot=None,
+               text_char_emb=None,
+               text_word_emb=None,
+               padding=None,
+               is_training=True):
+    """Applies Ithaca model on the inputs."""
+
+    if text_char is not None and padding is None:
+      padding = jnp.where(text_char > 0, 1, 0)
+    elif text_char_onehot is not None and padding is None:
+      padding = jnp.where(text_char_onehot.argmax(-1) > 0, 1, 0)
+    padding_mask = padding[..., jnp.newaxis]
+    text_len = jnp.sum(padding, 1)
+
+    if self.posemb_combine_type == 'add':
+      posemb_dim = None
+    elif self.posemb_combine_type == 'concat':
+      posemb_dim = self.word_char_emb_dim
+    else:
+      raise ValueError('Wrong feature_combine_type value.')
+
+    # Character embeddings
+    if text_char is not None:
+      x = self.text_char_emb(text_char)
+    elif text_char_onehot is not None:
+      x = self.text_char_emb.attend(text_char_onehot)
+    elif text_char_emb is not None:
+      x = text_char_emb
+    else:
+      raise ValueError('Wrong inputs.')
+
+    # Word embeddings
+    if text_word is not None:
+      text_word_emb_x = self.text_word_emb(text_word)
+    elif text_word_onehot is not None:
+      text_word_emb_x = self.text_word_emb.attend(text_word_onehot)
+    elif text_word_emb is not None:
+      text_word_emb_x = text_word_emb
+    else:
+      raise ValueError('Wrong inputs.')
+
+    if self.feature_combine_type == 'add':
+      x = x + text_word_emb_x
+    elif self.feature_combine_type == 'concat':
+      x = jax.lax.concatenate([x, text_word_emb_x], 2)
+    else:
+      raise ValueError('Wrong feature_combine_type value.')
+
+    # Positional embeddings
+    pe_init = common_layers.sinusoidal_init(
+        max_len=self.max_len) if self.learn_pos_emb else None
+    x = common_layers.AddPositionEmbs(
+        posemb_dim=posemb_dim,
+        posemb_init=pe_init,
+        max_len=self.max_len,
+        combine_type=self.posemb_combine_type,
+        name='posembed_input',
+    )(
+        x)
+    x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not is_training)
+
+    # Set floating point
+    if self.use_bfloat16:
+      x = x.astype(jnp.bfloat16)
+      dtype = jnp.bfloat16
+    else:
+      dtype = jnp.float32
+
+    if self.model_type == 'bigbird':
+      model_block = bigbird.BigBirdBlock
+    else:
+      raise ValueError('Wrong model type specified.')
+
+    for lyr in range(self.num_layers):
+      x = model_block(
+          qkv_dim=self.qkv_dim,
+          mlp_dim=self.mlp_dim,
+          num_heads=self.num_heads,
+          dtype=dtype,
+          causal_mask=self.causal_mask,
+          dropout_rate=self.dropout_rate,
+          attention_dropout_rate=self.attention_dropout_rate,
+          deterministic=not is_training,
+          activation_fn=self.activation_fn,
+          connectivity_seed=lyr,
+          name=f'encoderblock_{lyr}',
+      )(
+          x,
+          padding_mask=padding_mask,
+      )
+    x = common_layers.LayerNorm(dtype=dtype, name='encoder_norm')(x)
+    torso_output = x
+
+    # Bert logits
+    if self.use_output_mlp:
+      x_mask = common_layers.MlpBlock(
+          out_dim=self.word_char_emb_dim,
+          mlp_dim=self.emb_dim,
+          dtype=dtype,
+          out_dropout=False,
+          dropout_rate=self.dropout_rate,
+          deterministic=not is_training,
+          activation_fn=self.activation_fn)(
+              x)
+    else:
+      x_mask = nn.Dense(self.word_char_emb_dim)(x)
+
+    char_embeddings = self.text_char_emb.embedding
+    char_embeddings = nn.Dropout(rate=self.dropout_rate)(
+        char_embeddings, deterministic=not is_training)
+    logits_mask = jnp.matmul(x_mask, jnp.transpose(char_embeddings))
+
+    # Next sentence prediction
+    if self.use_output_mlp:
+      logits_nsp = common_layers.MlpBlock(
+          out_dim=2,
+          mlp_dim=self.emb_dim,
+          dtype=dtype,
+          out_dropout=False,
+          dropout_rate=self.dropout_rate,
+          deterministic=not is_training,
+          activation_fn=self.activation_fn)(
+              x)
+    else:
+      logits_nsp = nn.Dense(2)(x)
+
+    # Average over temporal dimension
+    if self.region_date_pooling == 'average':
+      x = jnp.multiply(padding_mask.astype(jnp.float32), x)
+      x = jnp.sum(x, 1) / text_len.astype(jnp.float32)[..., None]
+    elif self.region_date_pooling == 'sum':
+      x = jnp.multiply(padding_mask.astype(jnp.float32), x)
+      x = jnp.sum(x, 1)
+    elif self.region_date_pooling == 'first':
+      x = x[:, 0, :]
+    else:
+      raise ValueError('Wrong pooling type specified.')
+
+    # Date pred
+    if self.output_date_dist:
+      output_date_dim = self.output_date
+    else:
+      output_date_dim = 1
+
+    if self.use_output_mlp:
+      pred_date = common_layers.MlpBlock(
+          out_dim=output_date_dim,
+          mlp_dim=self.emb_dim,
+          dtype=dtype,
+          out_dropout=False,
+          dropout_rate=self.dropout_rate,
+          deterministic=not is_training,
+          activation_fn=self.activation_fn)(
+              x)
+    else:
+      pred_date = nn.Dense(output_date_dim)(x)
+
+    # Region logits
+    if self.use_output_mlp:
+      logits_subregion = common_layers.MlpBlock(
+          out_dim=self.output_subregions,
+          mlp_dim=self.emb_dim,
+          dtype=dtype,
+          out_dropout=False,
+          dropout_rate=self.dropout_rate,
+          deterministic=not is_training,
+          activation_fn=self.activation_fn)(
+              x)
+    else:
+      logits_subregion = nn.Dense(self.output_subregions)(x)
+
+    outputs = (pred_date, logits_subregion, logits_mask, logits_nsp)
+    if self.output_return_emb:
+      return outputs, torso_output
+    else:
+      return outputs

ithaca/util/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

ithaca/util/alphabet.py

@@ -0,0 +1,171 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Alphabet classes."""
+
+import re
+
+import numpy as np
+
+
+class Alphabet:
+  """Generic alphabet class."""
+
+  def __init__(self,
+               alphabet,
+               numerals='0',
+               punctuation='.',
+               space=' ',
+               missing='-',
+               pad='#',
+               unk='^',
+               sos='<',
+               sog='[',
+               eog=']',
+               wordlist_file=None,
+               wordlist_size=100000):
+    self.alphabet = list(alphabet)  # alph
+    self.numerals = list(numerals)  # num
+    self.punctuation = list(punctuation)  # punt
+    self.space = space  # spacing
+    self.missing = missing  # missing char
+    self.pad = pad  # padding (spaces to right of string)
+    self.unk = unk  # unknown char
+    self.sos = sos  # start of sentence
+    self.sog = sog  # start of guess
+    self.eog = eog  # end of guess
+
+    # Define wordlist mapping
+    idx2word = [self.pad, self.sos, self.unk]
+    if wordlist_file:
+      idx2word += [
+          w_c.split(';')[0]
+          for w_c in wordlist_file.read().strip().split('\n')[:wordlist_size]
+      ]
+    self.idx2word = np.array(idx2word)
+    self.word2idx = {self.idx2word[i]: i for i in range(len(self.idx2word))}
+
+    # Define vocab mapping
+    self.idx2char = np.array(
+        [self.pad, self.sos, self.unk, self.space, self.missing] +
+        self.alphabet + self.numerals + self.punctuation)
+    self.char2idx = {self.idx2char[i]: i for i in range(len(self.idx2char))}
+
+    # Define special character indices
+    self.pad_idx = self.char2idx[pad]
+    self.sos_idx = self.char2idx[sos]
+    self.unk_idx = self.char2idx[unk]
+    self.alphabet_start_idx = self.char2idx[self.alphabet[0]]
+    self.alphabet_end_idx = self.char2idx[self.numerals[-1]]
+
+  def filter(self, t):
+    return t
+
+  def size_char(self):
+    return len(self.idx2char)
+
+  def size_word(self):
+    return len(self.idx2word)
+
+
+class GreekAlphabet(Alphabet):
+  """Greek alphabet class."""
+
+  def __init__(self, wordlist_file=None, wordlist_size=100000):
+    greek_alphabet = 'αβγδεζηθικλμνξοπρςστυφχψωϙϛ'
+
+    super().__init__(
+        alphabet=greek_alphabet,
+        wordlist_file=wordlist_file,
+        wordlist_size=wordlist_size)
+    self.tonos_to_oxia = {
+        # tonos  : #oxia
+        u'\u0386': u'\u1FBB',  # capital letter alpha
+        u'\u0388': u'\u1FC9',  # capital letter epsilon
+        u'\u0389': u'\u1FCB',  # capital letter eta
+        u'\u038C': u'\u1FF9',  # capital letter omicron
+        u'\u038A': u'\u1FDB',  # capital letter iota
+        u'\u038E': u'\u1FF9',  # capital letter upsilon
+        u'\u038F': u'\u1FFB',  # capital letter omega
+        u'\u03AC': u'\u1F71',  # small letter alpha
+        u'\u03AD': u'\u1F73',  # small letter epsilon
+        u'\u03AE': u'\u1F75',  # small letter eta
+        u'\u0390': u'\u1FD3',  # small letter iota with dialytika and tonos/oxia
+        u'\u03AF': u'\u1F77',  # small letter iota
+        u'\u03CC': u'\u1F79',  # small letter omicron
+        u'\u03B0': u'\u1FE3',
+        # small letter upsilon with dialytika and tonos/oxia
+        u'\u03CD': u'\u1F7B',  # small letter upsilon
+        u'\u03CE': u'\u1F7D'  # small letter omega
+    }
+    self.oxia_to_tonos = {v: k for k, v in self.tonos_to_oxia.items()}
+
+  def filter(self, t):  # override previous filter function
+    # lowercase
+    t = t.lower()
+
+    # replace dot below
+    t = t.replace(u'\u0323', '')
+
+    # replace perispomeni
+    t = t.replace(u'\u0342', '')
+    t = t.replace(u'\u02C9', '')
+
+    # replace ending sigma
+    t = re.sub(r'([\w\[\]])σ(?![\[\]])(\b)', r'\1ς\2', t)
+
+    # replace oxia with tonos
+    for oxia, tonos in self.oxia_to_tonos.items():
+      t = t.replace(oxia, tonos)
+
+    # replace h
+    h_patterns = {
+        # input: #target
+        'ε': 'ἑ',
+        'ὲ': 'ἓ',
+        'έ': 'ἕ',
+        'α': 'ἁ',
+        'ὰ': 'ἃ',
+        'ά': 'ἅ',
+        'ᾶ': 'ἇ',
+        'ι': 'ἱ',
+        'ὶ': 'ἳ',
+        'ί': 'ἵ',
+        'ῖ': 'ἷ',
+        'ο': 'ὁ',
+        'ό': 'ὅ',
+        'ὸ': 'ὃ',
+        'υ': 'ὑ',
+        'ὺ': 'ὓ',
+        'ύ': 'ὕ',
+        'ῦ': 'ὗ',
+        'ὴ': 'ἣ',
+        'η': 'ἡ',
+        'ή': 'ἥ',
+        'ῆ': 'ἧ',
+        'ὼ': 'ὣ',
+        'ώ': 'ὥ',
+        'ω': 'ὡ',
+        'ῶ': 'ὧ'
+    }
+
+    # iterate by keys
+    for h_in, h_tar in h_patterns.items():
+      # look up and replace h[ and h]
+      t = re.sub(r'ℎ(\[?){}'.format(h_in), r'\1{}'.format(h_tar), t)
+      t = re.sub(r'ℎ(\]?){}'.format(h_in), r'{}\1'.format(h_tar), t)
+
+    # any h left is an ἡ
+    t = re.sub(r'(\[?)ℎ(\]?)', r'\1ἡ\2', t)
+
+    return t

ithaca/util/dates.py

@@ -0,0 +1,67 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Date processing functions."""
+import numpy as np
+
+
+def date_num_bins(date_min, date_max, date_interval, unknown_bin=True):
+  num_bins = (date_max - date_min - 1) // date_interval
+  if unknown_bin:
+    num_bins += 1  # +1 for unk
+  return num_bins
+
+
+def date_to_bin(date_cur, date_min, date_max, date_interval, date_bins):
+  if date_cur >= date_min and date_cur < date_max:
+    date_bin = np.digitize(
+        date_cur,
+        list(range(date_min + date_interval, date_max, date_interval)))
+  else:
+    date_bin = date_bins - 1
+  return date_bin
+
+
+def bin_to_date(date_cur_bin, date_min, date_interval):
+  return date_min + date_cur_bin * date_interval + date_interval // 2
+
+
+def date_range_to_dist(date_min_cur,
+                       date_max_cur,
+                       date_min,
+                       date_max,
+                       date_interval,
+                       date_bins,
+                       return_logits=True):
+  """Converts a date range to a uniform distribution."""
+  dist = np.zeros(date_bins)
+
+  if (date_min_cur and date_max_cur and date_min_cur >= date_min and
+      date_max_cur < date_max and date_min_cur <= date_max_cur):
+    date_min_cur_bin = date_to_bin(date_min_cur, date_min, date_max,
+                                   date_interval, date_bins)
+    date_max_cur_bin = date_to_bin(date_max_cur, date_min, date_max,
+                                   date_interval, date_bins)
+  else:
+    date_min_cur_bin = date_bins - 1
+    date_max_cur_bin = date_bins - 1
+
+  date_bins_cur = date_max_cur_bin - date_min_cur_bin + 1
+  dist[date_min_cur_bin:date_max_cur_bin + 1] = 1. / date_bins_cur
+
+  if return_logits:
+    eps = 1e-6
+    dist = np.clip(dist, eps, 1. - eps)
+    dist = np.log(dist)
+
+  return dist

ithaca/util/eval.py

@@ -0,0 +1,562 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Eval utils."""
+
+from typing import List, NamedTuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from .text import idx_to_text
+from .text import text_to_idx
+from .text import text_to_word_idx
+import tqdm
+
+
+def date_loss_l1(pred, target_min, target_max):
+  """L1 loss function for dates."""
+  loss = 0.
+  loss += np.abs(pred - target_min) * np.less(pred, target_min).astype(
+      pred.dtype)
+  loss += np.abs(pred - target_max) * np.greater(pred, target_max).astype(
+      pred.dtype)
+  return loss
+
+
+def grad_to_saliency_char(gradient_char, text_char_onehot, text_len, alphabet):
+  """Generates saliency map."""
+  saliency_char = np.linalg.norm(gradient_char, axis=2)[0, :text_len[0]]
+
+  text_char = np.array(text_char_onehot).argmax(axis=-1)
+  idx_mask = np.logical_or(
+      text_char[0, :text_len[0]] > alphabet.alphabet_end_idx,
+      text_char[0, :text_len[0]] < alphabet.alphabet_start_idx)
+  idx_unmask = np.logical_not(idx_mask)
+
+  saliency_char_tmp = saliency_char.copy()
+  saliency_char_tmp[idx_mask] = 0.
+  if idx_unmask.any():
+    saliency_char_tmp[idx_unmask] = (saliency_char[idx_unmask] -
+                                     saliency_char[idx_unmask].min()) / (
+                                         saliency_char[idx_unmask].max() -
+                                         saliency_char[idx_unmask].min() + 1e-8)
+  return saliency_char_tmp
+
+
+def grad_to_saliency_word(gradient_word, text_word_onehot, text_len, alphabet):
+  """Generates saliency map."""
+  saliency_word = np.linalg.norm(gradient_word, axis=2)[0, :text_len[0]]
+  text_word = np.array(text_word_onehot).argmax(axis=-1)
+
+  saliency_word = saliency_word.copy()
+  start_idx = None
+  for i in range(text_len[0]):
+    if text_word[0, i] == alphabet.unk_idx:
+      if start_idx is not None:
+        saliency_word[start_idx:i] = np.sum(saliency_word[start_idx:i])
+      start_idx = None
+    elif start_idx is None:
+      start_idx = i
+
+  idx_mask = text_word[0, :text_len[0]] == alphabet.unk_idx
+  idx_unmask = np.logical_not(idx_mask)
+  saliency_word_tmp = saliency_word.copy()
+  saliency_word_tmp[idx_mask] = 0.
+  if idx_unmask.any():
+    saliency_word_tmp[idx_unmask] = (
+        saliency_word[idx_unmask] - saliency_word[idx_unmask].min())
+    saliency_word_tmp[idx_unmask] = saliency_word_tmp[idx_unmask] / (
+        saliency_word[idx_unmask].max() - saliency_word[idx_unmask].min() +
+        1e-8)
+  return saliency_word_tmp
+
+
+def softmax(x, axis=-1):
+  """Compute softmax values for each sets of scores in x."""
+  unnormalized = np.exp(x - x.max(axis, keepdims=True))
+  return unnormalized / unnormalized.sum(axis, keepdims=True)
+
+
+def log_softmax(x, axis=-1):
+  """Log-Softmax function."""
+  shifted = x - x.max(axis, keepdims=True)
+  return shifted - np.log(np.sum(np.exp(shifted), axis, keepdims=True))
+
+
+def nucleus_sample_inner(logits, top_p=0.95, temp=1.0):
+  """Samples from the most likely tokens whose probability sums to top_p."""
+  sorted_logits = np.sort(logits)
+  sorted_probs = softmax(sorted_logits)
+  threshold_idx = np.argmax(np.cumsum(sorted_probs, -1) >= 1 - top_p)
+  threshold_largest_logits = sorted_logits[..., [threshold_idx]]
+  assert threshold_largest_logits.shape == logits.shape[:-1] + (1,)
+  mask = logits >= threshold_largest_logits
+  logits += (1 - mask) * -1e12  # Set unused logits to -inf.
+  logits /= np.maximum(temp, 1e-12)
+  return logits
+
+
+class BeamEntry(NamedTuple):
+  text_pred: str
+  mask_idx: int
+  pred_len: int
+  pred_logprob: float
+
+
+def beam_search_batch_2d(forward,
+                         alphabet,
+                         text_pred,
+                         mask_idx,
+                         rng=None,
+                         beam_width=20,
+                         temperature=1.,
+                         nucleus=False,
+                         nucleus_top_p=0.8,
+                         display_progress=False) -> List[BeamEntry]:
+  """Non-sequential beam search."""
+
+  beam = [BeamEntry(text_pred, mask_idx, 0, 0.)]
+  beam_top = {}
+
+  text_len = len(text_pred.rstrip(alphabet.pad))
+
+  # Initialise tqdm bar
+  if display_progress:
+    pbar = tqdm.tqdm(total=len(mask_idx))
+
+  while beam:
+    beam_tmp = []
+    beam_batch = []
+
+    text_chars = []
+    text_words = []
+
+    for text_pred, mask_idx, pred_len, pred_logprob in beam:
+
+      mask_idx = mask_idx.copy()  # pytype: disable=attribute-error  # strict_namedtuple_checks
+      text_char = text_to_idx(text_pred, alphabet).reshape(1, -1)
+      text_word = text_to_word_idx(text_pred, alphabet).reshape(1, -1)
+      text_chars.append(text_char)
+      text_words.append(text_word)
+      beam_batch.append(BeamEntry(text_pred, mask_idx, pred_len, pred_logprob))
+    text_chars = np.vstack(text_chars)
+    text_words = np.vstack(text_words)
+
+    _, _, mask_logits, _ = forward(
+        text_char=text_chars,
+        text_word=text_words,
+        text_char_onehot=None,
+        text_word_onehot=None,
+        rngs={'dropout': rng},
+        is_training=False)
+    mask_logits = mask_logits / temperature
+    mask_logits = np.array(mask_logits)
+
+    for batch_i in range(mask_logits.shape[0]):
+      text_pred, mask_idx, pred_len, pred_logprob = beam_batch[batch_i]
+      mask_logprob = log_softmax(mask_logits)[batch_i, :text_len]
+      mask_pred = softmax(mask_logits)[batch_i, :text_len]
+      mask_pred_argmax = np.dstack(
+          np.unravel_index(np.argsort(-mask_pred.ravel()), mask_pred.shape))[0]
+
+      # Keep only predictions for mask
+      for i in range(mask_pred_argmax.shape[0]):
+        if (mask_pred_argmax[i][0] in mask_idx and  # pytype: disable=unsupported-operands  # strict_namedtuple_checks
+            (mask_pred_argmax[i][1] == alphabet.char2idx[alphabet.space] or
+             (mask_pred_argmax[i][1] >= alphabet.alphabet_start_idx and
+              mask_pred_argmax[i][1] <=
+              alphabet.char2idx[alphabet.punctuation[-1]]))):
+          text_char_i = text_chars.copy()
+          text_char_i[batch_i, mask_pred_argmax[i][0]] = mask_pred_argmax[i][1]
+          text_pred_i = idx_to_text(
+              text_char_i[batch_i], alphabet, strip_sos=False, strip_pad=False)
+
+          mask_idx_i = mask_idx.copy()  # pytype: disable=attribute-error  # strict_namedtuple_checks
+          mask_idx_i.remove(mask_pred_argmax[i][0])
+
+          if nucleus:
+            mask_logits_i = mask_logits[batch_i, mask_pred_argmax[i][0]]
+            mask_logits_i = nucleus_sample_inner(mask_logits_i, nucleus_top_p)
+            mask_logprob_i = log_softmax(mask_logits_i)
+
+            # Skip expanding the beam if logprob too small
+            if mask_logits_i[mask_pred_argmax[i][1]] < -1e12:
+              continue
+
+            pred_logprob_i = pred_logprob + mask_logprob_i[mask_pred_argmax[i]
+                                                           [1]]
+          else:
+            pred_logprob_i = pred_logprob + mask_logprob[mask_pred_argmax[i][0],
+                                                         mask_pred_argmax[i][1]]
+
+          if not mask_idx_i:
+            if (text_pred_i
+                not in beam_top) or (text_pred_i in beam_top and
+                                     beam_top[text_pred_i][3] > pred_logprob_i):
+              beam_top[text_pred_i] = BeamEntry(text_pred_i, mask_idx_i,
+                                                pred_len + 1, pred_logprob_i)
+          else:
+            beam_tmp.append(
+                BeamEntry(text_pred_i, mask_idx_i, pred_len + 1,
+                          pred_logprob_i))
+
+    # order all candidates by score
+    beam_tmp_kv = {}
+    for text_pred, mask_idx, pred_len, pred_logprob in beam_tmp:
+      if (text_pred not in beam_tmp_kv) or (
+          text_pred in beam_tmp_kv and
+          beam_tmp_kv[text_pred].pred_logprob > pred_logprob):
+        beam_tmp_kv[text_pred] = BeamEntry(text_pred, mask_idx, pred_len,
+                                           pred_logprob)
+    beam_tmp = sorted(
+        beam_tmp_kv.values(),
+        key=lambda entry: entry.pred_logprob,
+        reverse=True)
+
+    # select k best
+    beam = beam_tmp[:beam_width]
+
+    # update progress bar
+    if display_progress:
+      pbar.update(1)
+
+  # order all candidates by score
+  return sorted(
+      beam_top.values(), key=lambda entry: entry.pred_logprob,
+      reverse=True)[:beam_width]
+
+
+def beam_search_batch_1d(forward,
+                         alphabet,
+                         text_pred,
+                         mask_idx,
+                         rng=None,
+                         beam_width=20,
+                         temperature=1.,
+                         nucleus=False,
+                         nucleus_top_p=0.8,
+                         display_progress=False) -> List[BeamEntry]:
+  """Sequential beam search."""
+
+  beam = [BeamEntry(text_pred, mask_idx, 0, 0.)]
+  beam_top = {}
+
+  # Initialise tqdm bar
+  if display_progress:
+    pbar = tqdm.tqdm(total=len(mask_idx))
+
+  while beam:
+    beam_tmp = []
+    beam_batch = []
+
+    text_chars = []
+    text_words = []
+
+    for text_pred, mask_idx, pred_len, pred_logprob in beam:
+
+      mask_idx = mask_idx.copy()  # pytype: disable=attribute-error  # strict_namedtuple_checks
+      text_char = text_to_idx(text_pred, alphabet).reshape(1, -1)
+      text_word = text_to_word_idx(text_pred, alphabet).reshape(1, -1)
+      text_chars.append(text_char)
+      text_words.append(text_word)
+      beam_batch.append(BeamEntry(text_pred, mask_idx, pred_len, pred_logprob))
+    text_chars = np.vstack(text_chars)
+    text_words = np.vstack(text_words)
+
+    _, _, mask_logits, _ = forward(
+        text_char=text_chars,
+        text_word=text_words,
+        text_char_onehot=None,
+        text_word_onehot=None,
+        rngs={'dropout': rng},
+        is_training=False)
+    mask_logits = mask_logits / temperature
+    mask_logits = np.array(mask_logits)
+
+    for batch_i in range(mask_logits.shape[0]):
+      text_pred, mask_idx, pred_len, pred_logprob = beam_batch[batch_i]
+
+      mask_logits_i = mask_logits[batch_i, mask_idx[0]]  # pytype: disable=unsupported-operands  # strict_namedtuple_checks
+      if nucleus:
+        mask_logits_i = nucleus_sample_inner(mask_logits_i, nucleus_top_p)
+
+      mask_logprob = log_softmax(mask_logits_i)
+
+      # Keep only predictions for mask
+      alphabet_chars = [alphabet.char2idx[alphabet.space]]
+      alphabet_chars += list(
+          range(alphabet.alphabet_start_idx,
+                alphabet.char2idx[alphabet.punctuation[-1]]))
+      for char_i in alphabet_chars:
+        # Skip expanding the beam if logprob too small
+        if nucleus and mask_logits_i[char_i] < -1e12:
+          continue
+
+        text_char_i = text_chars.copy()
+        text_char_i[batch_i, mask_idx[0]] = char_i  # pytype: disable=unsupported-operands  # strict_namedtuple_checks
+
+        text_pred_i = idx_to_text(
+            text_char_i[batch_i], alphabet, strip_sos=False, strip_pad=False)
+
+        mask_idx_i = mask_idx.copy()  # pytype: disable=attribute-error  # strict_namedtuple_checks
+        mask_idx_i.pop(0)
+        pred_logprob_i = pred_logprob + mask_logprob[char_i]
+
+        if not mask_idx_i:
+          if (text_pred_i
+              not in beam_top) or (text_pred_i in beam_top and
+                                   beam_top[text_pred_i][3] > pred_logprob_i):
+            beam_top[text_pred_i] = BeamEntry(text_pred_i, mask_idx_i,
+                                              pred_len + 1, pred_logprob_i)
+        else:
+          beam_tmp.append(
+              BeamEntry(text_pred_i, mask_idx_i, pred_len + 1, pred_logprob_i))
+
+    # order all candidates by score
+    beam_tmp_kv = {}
+    for text_pred, mask_idx, pred_len, pred_logprob in beam_tmp:
+      if (text_pred
+          not in beam_tmp_kv) or (text_pred in beam_tmp_kv and
+                                  beam_tmp_kv[text_pred][3] > pred_logprob):
+        beam_tmp_kv[text_pred] = BeamEntry(text_pred, mask_idx, pred_len,
+                                           pred_logprob)
+    beam_tmp = sorted(
+        beam_tmp_kv.values(),
+        key=lambda entry: entry.pred_logprob,
+        reverse=True)
+
+    # select k best
+    beam = beam_tmp[:beam_width]
+
+    # update progress bar
+    if display_progress:
+      pbar.update(1)
+
+  # order all candidates by score
+  return sorted(
+      beam_top.values(), key=lambda entry: entry.pred_logprob,
+      reverse=True)[:beam_width]
+
+
+def saliency_loss_subregion(forward,
+                            text_char_emb,
+                            text_word_emb,
+                            padding,
+                            rng,
+                            subregion=None):
+  """Saliency map for subregion."""
+
+  _, subregion_logits, _, _ = forward(
+      text_char_emb=text_char_emb,
+      text_word_emb=text_word_emb,
+      padding=padding,
+      rngs={'dropout': rng},
+      is_training=False)
+  if subregion is None:
+    subregion = subregion_logits.argmax(axis=-1)[0]
+  return subregion_logits[0, subregion]
+
+
+def saliency_loss_date(forward, text_char_emb, text_word_emb, padding, rng):
+  """saliency_loss_date."""
+
+  date_pred, _, _, _ = forward(
+      text_char_emb=text_char_emb,
+      text_word_emb=text_word_emb,
+      padding=padding,
+      rngs={'dropout': rng},
+      is_training=False)
+
+  date_pred_argmax = date_pred.argmax(axis=-1)
+  return date_pred[0, date_pred_argmax[0]]
+
+
+def predicted_dates(date_pred_probs, date_min, date_max, date_interval):
+  """Returns mode and mean prediction."""
+  date_years = np.arange(date_min + date_interval / 2,
+                         date_max + date_interval / 2, date_interval)
+
+  # Compute mode:
+  date_pred_argmax = (
+      date_pred_probs.argmax() * date_interval + date_min + date_interval // 2)
+
+  # Compute mean:
+  date_pred_avg = np.dot(date_pred_probs, date_years)
+
+  return date_pred_argmax, date_pred_avg
+
+
+def compute_attribution_saliency_maps(text_char,
+                                      text_word,
+                                      text_len,
+                                      padding,
+                                      forward,
+                                      params,
+                                      rng,
+                                      alphabet,
+                                      vocab_char_size,
+                                      vocab_word_size,
+                                      subregion_loss_kwargs=None):
+  """Compute saliency maps for subregions and dates."""
+
+  if subregion_loss_kwargs is None:
+    subregion_loss_kwargs = {}
+
+  # Get saliency gradients
+  dtype = params['params']['char_embeddings']['embedding'].dtype
+  text_char_onehot = jax.nn.one_hot(text_char, vocab_char_size).astype(dtype)
+  text_word_onehot = jax.nn.one_hot(text_word, vocab_word_size).astype(dtype)
+  text_char_emb = jnp.matmul(text_char_onehot,
+                             params['params']['char_embeddings']['embedding'])
+  text_word_emb = jnp.matmul(text_word_onehot,
+                             params['params']['word_embeddings']['embedding'])
+  gradient_subregion_char, gradient_subregion_word = jax.grad(
+      saliency_loss_subregion, (1, 2))(
+          forward,
+          text_char_emb,
+          text_word_emb,
+          padding,
+          rng=rng,
+          **subregion_loss_kwargs)
+  gradient_date_char, gradient_date_word = jax.grad(saliency_loss_date, (1, 2))(
+      forward, text_char_emb, text_word_emb, padding=padding, rng=rng)
+
+  # Generate saliency maps for subregions
+  input_grad_subregion_char = np.multiply(gradient_subregion_char,
+                                          text_char_emb)  # grad x input
+  input_grad_subregion_word = np.multiply(gradient_subregion_word,
+                                          text_word_emb)
+  grad_char = grad_to_saliency_char(
+      input_grad_subregion_char,
+      text_char_onehot,
+      text_len=text_len,
+      alphabet=alphabet)
+  grad_word = grad_to_saliency_word(
+      input_grad_subregion_word,
+      text_word_onehot,
+      text_len=text_len,
+      alphabet=alphabet)
+  subregion_saliency = np.clip(grad_char + grad_word, 0, 1)
+
+  # Generate saliency maps for dates
+  input_grad_date_char = np.multiply(gradient_date_char,
+                                     text_char_emb)  # grad x input
+  input_grad_date_word = np.multiply(gradient_date_word, text_word_emb)
+  grad_char = grad_to_saliency_char(
+      input_grad_date_char,
+      text_char_onehot,
+      text_len=text_len,
+      alphabet=alphabet)
+  grad_word = grad_to_saliency_word(
+      input_grad_date_word,
+      text_word_onehot,
+      text_len=text_len,
+      alphabet=alphabet)
+  date_saliency = np.clip(grad_char + grad_word, 0, 1)
+
+  return date_saliency, subregion_saliency
+
+
+def saliency_loss_mask(forward, text_char_emb, text_word_emb, padding, rng,
+                       char_pos, char_idx):
+  """Saliency map for mask."""
+
+  _, _, mask_logits, _ = forward(
+      text_char_emb=text_char_emb,
+      text_word_emb=text_word_emb,
+      text_char_onehot=None,
+      text_word_onehot=None,
+      padding=padding,
+      rngs={'dropout': rng},
+      is_training=False)
+  return mask_logits[0, char_pos, char_idx]
+
+
+class SequentialRestorationSaliencyResult(NamedTuple):
+  text: str  # predicted text string so far
+  pred_char_pos: int  # newly restored character's position
+  saliency_map: np.ndarray  # saliency map for the newly added character
+
+
+def sequential_restoration_saliency(text_str, text_len, forward, params,
+                                    alphabet, mask_idx, vocab_char_size,
+                                    vocab_word_size):
+  """Greedily, non-sequentially restores, producing per-step saliency maps."""
+  text_len = text_len[0] if not isinstance(text_len, int) else text_len
+  rng = jax.random.PRNGKey(0)  # dummy, no randomness in model
+  mask_idx = set(mask_idx)
+  while mask_idx:
+    text_char = text_to_idx(text_str, alphabet).reshape(1, -1)
+    padding = jnp.where(text_char > 0, 1, 0)
+    text_word = text_to_word_idx(text_str, alphabet).reshape(1, -1)
+
+    _, _, mask_logits, _ = forward(
+        text_char=text_char,
+        text_word=text_word,
+        text_char_onehot=None,
+        text_word_onehot=None,
+        rngs={'dropout': rng},
+        is_training=False)
+    mask_pred = jax.nn.softmax(mask_logits)[0, :text_len]
+    mask_pred_argmax = np.dstack(
+        np.unravel_index(np.argsort(-mask_pred.ravel()), mask_pred.shape))[0]
+
+    # Greedily, non-sequentially take the next highest probability prediction
+    # out of the characters that are to be restored
+    for i in range(mask_pred_argmax.shape[0]):
+      pred_char_pos, pred_char_idx = mask_pred_argmax[i]
+      if pred_char_pos in mask_idx:
+        break
+
+    # Update sequence
+    text_char[0, pred_char_pos] = pred_char_idx
+    text_str = idx_to_text(
+        text_char[0], alphabet, strip_sos=False, strip_pad=False)
+    mask_idx.remove(pred_char_pos)
+
+    # Gradients for saliency map
+    text_char_onehot = jax.nn.one_hot(text_char,
+                                      vocab_char_size).astype(jnp.float32)
+    text_word_onehot = jax.nn.one_hot(text_word,
+                                      vocab_word_size).astype(jnp.float32)
+
+    text_char_emb = jnp.matmul(text_char_onehot,
+                               params['params']['char_embeddings']['embedding'])
+    text_word_emb = jnp.matmul(text_word_onehot,
+                               params['params']['word_embeddings']['embedding'])
+
+    gradient_mask_char, gradient_mask_word = jax.grad(
+        saliency_loss_mask, (1, 2))(
+            forward,
+            text_char_emb,
+            text_word_emb,
+            padding,
+            rng=rng,
+            char_pos=pred_char_pos,
+            char_idx=pred_char_idx)
+
+    # Use gradient x input for visualizing saliency
+    input_grad_mask_char = np.multiply(gradient_mask_char, text_char_emb)
+    input_grad_mask_word = np.multiply(gradient_mask_word, text_word_emb)
+
+    # Return visualization-ready saliency maps
+    saliency_map = grad_to_saliency_char(
+        np.clip(input_grad_mask_char + input_grad_mask_word, 0, 1),
+        text_char_onehot, [text_len], alphabet)  # normalize, etc.
+    result_text = idx_to_text(text_char[0], alphabet, strip_sos=False)  # no pad
+
+    yield SequentialRestorationSaliencyResult(
+        text=result_text[1:],
+        pred_char_pos=pred_char_pos - 1,
+        saliency_map=saliency_map[1:])

ithaca/util/loss.py

@@ -0,0 +1,117 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Loss functions."""
+import chex
+from flax.deprecated import nn
+import jax
+import jax.numpy as jnp
+
+
+def smooth_labels(labels, num_classes, label_smoothing):
+  if not 0 <= label_smoothing < 1:
+    raise ValueError(
+        f"'label_smoothing is {label_smoothing} and should be in [0, 1)")
+  one = jax.lax.convert_element_type(1, labels.dtype)
+  label_smoothing = jax.lax.convert_element_type(label_smoothing,
+                                                 labels.dtype)
+  num_classes = jax.lax.convert_element_type(num_classes, labels.dtype)
+  return (one - label_smoothing) * labels + (label_smoothing / num_classes)
+
+
+def categorical_kl_divergence(p_logits, q_logits, temperature=1.):
+  """Compute the KL between two categorical distributions from their logits.
+
+  Args:
+    p_logits: unnormalized logits for the first distribution.
+    q_logits: unnormalized logits for the second distribution.
+    temperature: the temperature for the softmax distribution, defaults at 1.
+
+  Returns:
+    the kl divergence between the distributions.
+  """
+  chex.assert_type([p_logits, q_logits], float)
+
+  p_logits /= temperature
+  q_logits /= temperature
+
+  p = jax.nn.softmax(p_logits)
+  log_p = jax.nn.log_softmax(p_logits)
+  log_q = jax.nn.log_softmax(q_logits)
+  kl = jnp.sum(p * (log_p - log_q), axis=-1)
+  return jax.nn.relu(kl)  # Guard against numerical issues giving negative KL.
+
+
+def cross_entropy_label_smoothing_loss(logits,
+                                       labels,
+                                       mask=None,
+                                       label_smoothing=0.1):
+  """Cross entropy loss with label smoothing."""
+
+  num_classes = logits.shape[-1]
+  labels_onehot = jax.nn.one_hot(labels, num_classes, dtype=logits.dtype)
+  if label_smoothing > 0:
+    labels_onehot = smooth_labels(labels_onehot, num_classes, label_smoothing)
+
+  loss = -jnp.sum(labels_onehot * jax.nn.log_softmax(logits), axis=-1)
+  if mask is not None:
+    loss = jnp.multiply(loss, mask.astype(logits.dtype))
+  return loss
+
+
+@jax.vmap
+def cross_entropy_loss(logits, label):
+  logits = nn.log_softmax(logits)
+  return -logits[label]
+
+
+def cross_entropy_mask_loss(logits, label, mask):
+  nll = -nn.log_softmax(logits)[label]
+  loss = jnp.multiply(nll, mask.astype(logits.dtype))
+  return loss
+
+
+def date_loss_l2(pred,
+                 target_min,
+                 target_max,
+                 mask):
+  """L2 loss function for dates."""
+  pred = jnp.squeeze(pred, 0)
+
+  loss = 0.
+  loss += (pred - target_min)**2 * jnp.less(pred, target_min).astype(
+      pred.dtype)
+  loss += (pred - target_max)**2 * jnp.greater(pred, target_max).astype(
+      pred.dtype)
+
+  # Mask loss
+  loss = jnp.multiply(loss, mask.astype(loss.dtype))
+  return loss
+
+
+def date_loss_l1(pred,
+                 target_min,
+                 target_max,
+                 mask):
+  """L1 loss function for dates."""
+  pred = jnp.squeeze(pred, 0)
+
+  loss = 0.
+  loss += jnp.abs(pred - target_min) * jnp.less(pred, target_min).astype(
+      pred.dtype)
+  loss += jnp.abs(pred - target_max) * jnp.greater(pred, target_max).astype(
+      pred.dtype)
+
+  # Mask loss
+  loss = jnp.multiply(loss, mask.astype(loss.dtype))
+  return loss

ithaca/util/optim.py

@@ -0,0 +1,167 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Optimizer utilities."""
+
+from typing import Any, Callable, NamedTuple, Optional, Sequence, Tuple, Union
+import jax
+import jax.numpy as jnp
+
+
+def linear_weight(global_step, start, end):
+  """Linear weight increase."""
+  if end <= 0:
+    return 1.
+  t = jnp.maximum(0., global_step - start)
+  w = t / (end - start)
+  w = jnp.minimum(w, 1.)
+  return w
+
+
+def linear_warmup_and_sqrt_decay(global_step, max_lr, warmup_steps):
+  """Linear warmup and then an inverse square root decay of learning rate."""
+  linear_ratio = max_lr / warmup_steps
+  decay_ratio = jnp.power(warmup_steps * 1.0, 0.5) * max_lr
+  return jnp.minimum(linear_ratio * global_step,
+                     decay_ratio * jnp.power(global_step, -0.5))
+
+
+def create_learning_rate_scheduler(
+    factors='constant * linear_warmup * rsqrt_decay',
+    base_learning_rate=0.5,
+    warmup_steps=1000,
+    decay_factor=0.5,
+    steps_per_decay=20000,
+    steps_per_cycle=100000):
+  """Creates learning rate schedule.
+
+  Interprets factors in the factors string which can consist of:
+  * constant: interpreted as the constant value,
+  * linear_warmup: interpreted as linear warmup until warmup_steps,
+  * rsqrt_decay: divide by square root of max(step, warmup_steps)
+  * rsqrt_normalized_decay: divide by square root of max(step/warmup_steps, 1)
+  * decay_every: Every k steps decay the learning rate by decay_factor.
+  * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.
+
+  Args:
+    factors: string, factors separated by '*' that defines the schedule.
+    base_learning_rate: float, the starting constant for the lr schedule.
+    warmup_steps: int, how many steps to warm up for in the warmup schedule.
+    decay_factor: float, the amount to decay the learning rate by.
+    steps_per_decay: int, how often to decay the learning rate.
+    steps_per_cycle: int, steps per cycle when using cosine decay.
+
+  Returns:
+    a function learning_rate(step): float -> {'learning_rate': float}, the
+    step-dependent lr.
+  """
+  factors = [n.strip() for n in factors.split('*')]
+
+  def step_fn(step):
+    """Step to learning rate function."""
+    ret = 1.0
+    for name in factors:
+      if name == 'constant':
+        ret *= base_learning_rate
+      elif name == 'linear_warmup':
+        ret *= jnp.minimum(1.0, step / warmup_steps)
+      elif name == 'rsqrt_decay':
+        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
+      elif name == 'rsqrt_normalized_decay':
+        ret *= jnp.sqrt(warmup_steps)
+        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
+      elif name == 'decay_every':
+        ret *= (decay_factor**(step // steps_per_decay))
+      elif name == 'cosine_decay':
+        progress = jnp.maximum(0.0,
+                               (step - warmup_steps) / float(steps_per_cycle))
+        ret *= jnp.maximum(0.0,
+                           0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
+      else:
+        raise ValueError('Unknown factor %s.' % name)
+    return jnp.asarray(ret, dtype=jnp.float32)
+
+  return step_fn
+
+
+# pylint:disable=no-value-for-parameter
+OptState = NamedTuple  # Transformation states are (possibly empty) namedtuples.
+Params = Any  # Parameters are arbitrary nests of `jnp.ndarrays`.
+Updates = Params  # Gradient updates are of the same type as parameters.
+
+
+class GradientTransformation(NamedTuple):
+  """Optax transformations consists of a function pair: (initialise, update)."""
+  init: Callable[  # Function used to initialise the transformation's state.
+      [Params], Union[OptState, Sequence[OptState]]]
+  update: Callable[  # Function used to apply a transformation.
+      [Updates, OptState, Optional[Params]], Tuple[Updates, OptState]]
+
+
+class ClipByGlobalNormState(OptState):
+  """The `clip_by_global_norm` transformation is stateless."""
+
+
+def unitwise_norm(x):
+  """Computes norms of each output unit separately."""
+  if len(jnp.squeeze(x).shape) <= 1:  # Scalars and vectors
+    axis = None
+    keepdims = False
+  elif len(x.shape) in [2, 3]:  # Linear layers of shape IO or multihead linear
+    axis = 0
+    keepdims = True
+  elif len(x.shape) == 4:  # Conv kernels of shape HWIO
+    axis = [0, 1, 2,]
+    keepdims = True
+  else:
+    raise ValueError(f'Got a parameter with shape not in [1, 2, 3, 4]! {x}')
+  return jnp.sum(x ** 2, axis=axis, keepdims=keepdims) ** 0.5
+
+
+def unitwise_clip(g_norm, max_norm, grad):
+  """Applies gradient clipping unit-wise."""
+  trigger = g_norm < max_norm
+  # This little max(., 1e-6) is distinct from the normal eps and just prevents
+  # division by zero. It technically should be impossible to engage.
+  clipped_grad = grad * (max_norm / jnp.maximum(g_norm, 1e-6))
+  return jnp.where(trigger, grad, clipped_grad)
+
+
+def adaptive_grad_clip(clipping, eps=1e-3) -> GradientTransformation:
+  """Clip updates to be at most clipping * parameter_norm, unit-wise.
+
+  References:
+    [Brock, Smith, De, Simonyan 2021] High-Performance Large-Scale Image
+    Recognition Without Normalization. (https://arxiv.org/abs/2102.06171)
+
+  Args:
+    clipping: Maximum allowed ratio of update norm to parameter norm.
+    eps: epsilon term to prevent clipping of zero-initialized params.
+
+  Returns:
+    An (init_fn, update_fn) tuple.
+  """
+
+  def init_fn(_):
+    return ClipByGlobalNormState()
+
+  def update_fn(updates, state, params):
+    g_norm = jax.tree_map(unitwise_norm, updates)
+    p_norm = jax.tree_map(unitwise_norm, params)
+    # Maximum allowable norm
+    max_norm = jax.tree_map(lambda x: clipping * jnp.maximum(x, eps), p_norm)
+    # If grad norm > clipping * param_norm, rescale
+    updates = jax.tree_multimap(unitwise_clip, g_norm, max_norm, updates)
+    return updates, state
+
+  return GradientTransformation(init_fn, update_fn)

ithaca/util/region_names.py

@@ -0,0 +1,44 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Subregion mapping used to train the model.
+
+The subregion IDs originate from the I.PHI generator and may be subject to
+change in future versions of the PHI dataset.
+"""
+
+
+def load_region_maps(region_file):
+  """Extracts creates a map from PHI region id to a continuous region id."""
+  region_ids = []  # Used mainly for eval
+  region_ids_inv = {}  # Used in data loader
+  region_names_inv = {}  # Used in eval
+  for l in region_file.read().strip().split('\n'):
+    tok_name_id, _ = l.strip().split(';')  # second field is frequency, unused
+    region_name, region_id = tok_name_id.split('_')
+    region_name = region_name.strip()
+    region_id = int(region_id)
+    # Ignore unknown regions:
+    if ((region_name == 'Unknown Provenances' and region_id == 884) or
+        (region_name == 'unspecified subregion' and region_id == 885) or
+        (region_name == 'unspecified subregion' and region_id == 1439)):
+      continue
+    region_ids.append(region_id)
+    region_ids_inv[region_id] = len(region_ids_inv)
+    region_names_inv[len(region_names_inv)] = region_name
+
+  return {
+      'ids': region_ids,
+      'ids_inv': region_ids_inv,
+      'names_inv': region_names_inv
+  }

ithaca/util/text.py

@@ -0,0 +1,186 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Text processing functions."""
+
+import random
+import re
+import unicodedata
+
+import numpy as np
+
+
+def idx_to_text(idxs, alphabet, strip_sos=True, strip_pad=True):
+  """Converts a list of indices to a string."""
+  idxs = np.array(idxs)
+  out = ''
+  for i in range(idxs.size):
+    idx = idxs[i]
+    if strip_pad and idx == alphabet.pad_idx:
+      break
+    elif strip_sos and idx == alphabet.sos_idx:
+      pass
+    else:
+      out += alphabet.idx2char[idx]
+  return out
+
+
+def idx_to_text_batch(idxs, alphabet, lengths=None):
+  """Converts batched lists of indices to strings."""
+  b = []
+  for i in range(idxs.shape[0]):
+    idxs_i = idxs[i]
+    if lengths:
+      idxs_i = idxs_i[:lengths[i]]
+    b.append(idx_to_text(idxs_i, alphabet))
+  return b
+
+
+def random_mask_span(t, geometric_p=0.2, limit_chars=None):
+  """Masks a span of sequential words."""
+
+  # Obtain span indexes (indlusive)
+  span_idx = [(ele.start(), ele.end()) for ele in re.finditer(r'[\w\s]+', t)]
+  if not span_idx:
+    return []
+
+  # Select a span to mask
+  span_start, span_end = random.choice(span_idx)
+
+  # Sample a random span length using a geomteric distribution
+  if geometric_p and limit_chars:
+    span_len = np.clip(
+        np.random.geometric(geometric_p),
+        1, min(limit_chars, span_end - span_start))
+  elif geometric_p:
+    span_len = np.clip(
+        np.random.geometric(geometric_p),
+        1, span_end - span_start)
+  elif limit_chars:
+    span_len = min(limit_chars, span_end - span_start)
+  else:
+    raise ValueError('geometric_p or limit_chars should be set.')
+
+  # Pick a random start index
+  span_start = np.random.randint(span_start, span_end - span_len + 1)
+  assert span_start + span_len <= span_end
+
+  # Clip to limit chars
+  if limit_chars is not None and span_len >= limit_chars:
+    span_len = limit_chars
+
+  # Create mask indices
+  mask_idx = list(range(span_start, span_start + span_len))
+
+  return mask_idx
+
+
+def random_sentence_swap(sentences, p):
+  """Swaps sentences with probability p."""
+
+  def swap_sentence(s):
+    idx_1 = random.randint(0, len(s) - 1)
+    idx_2 = idx_1
+    counter = 0
+
+    while idx_2 == idx_1:
+      idx_2 = random.randint(0, len(s) - 1)
+      counter += 1
+      if counter > 3:
+        return s
+
+    s[idx_1], s[idx_2] = s[idx_2], s[idx_1]
+    return s
+
+  new_sentences = sentences.copy()
+  n = int(p * len(sentences))
+  for _ in range(n):
+    new_sentences = swap_sentence(new_sentences)
+
+  return new_sentences
+
+
+def random_word_delete(sentence, p):
+  """Deletes a word from a sentence with probability p."""
+
+  words = sentence.split(' ')
+
+  # Return if one word.
+  if len(words) == 1:
+    return words[0]
+
+  # Randomly delete words.
+  new_words = []
+  for word in words:
+    if random.uniform(0, 1) > p:
+      new_words.append(word)
+
+  # If all words are removed return one.
+  if not new_words:
+    rand_int = random.randint(0, len(words) - 1)
+    return words[rand_int]
+
+  sentence = ' '.join(new_words)
+
+  return sentence
+
+
+def random_word_swap(sentence, p):
+  """Swaps words from a sentence with probability p."""
+
+  def swap_word(new_words):
+    idx_1 = random.randint(0, len(new_words) - 1)
+    idx_2 = idx_1
+    counter = 0
+
+    while idx_2 == idx_1:
+      idx_2 = random.randint(0, len(new_words) - 1)
+      counter += 1
+
+      if counter > 3:
+        return new_words
+
+    new_words[idx_1], new_words[idx_2] = new_words[idx_2], new_words[idx_1]
+    return new_words
+
+  words = sentence.split(' ')
+
+  new_words = words.copy()
+  n = int(p * len(words))
+  for _ in range(n):
+    new_words = swap_word(new_words)
+
+  sentence = ' '.join(new_words)
+
+  return sentence
+
+
+def strip_accents(s):
+  return ''.join(
+      c for c in unicodedata.normalize('NFD', s)
+      if unicodedata.category(c) != 'Mn')
+
+
+def text_to_idx(t, alphabet):
+  """Converts a string to character indices."""
+  return np.array([alphabet.char2idx[c] for c in t], dtype=np.int32)
+
+
+def text_to_word_idx(t, alphabet):
+  """Converts a string to word indices."""
+  out = np.full(len(t), alphabet.word2idx[alphabet.unk], dtype=np.int32)
+  for m in re.finditer(r'\w+', t):
+    if m.group() in alphabet.word2idx:
+      out[m.start():m.end()] = alphabet.word2idx[m.group()]
+  return out
+

requirements.txt

@@ -0,0 +1,8 @@
+absl-py==0.13.0
+chex==0.0.8
+flax==0.3.6
+dm-haiku==0.0.5
+jax==0.2.21
+ml-collections==0.1.1
+numpy>=1.18.0
+tqdm>=4.62.2

setup.py

@@ -0,0 +1,52 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Setup module for Ithaca.
+
+Only installs the inference components.
+
+See README.md for more details.
+"""
+
+import pathlib
+import setuptools
+
+here = pathlib.Path(__file__).parent.resolve()
+long_description = (here / 'README.md').read_text(encoding='utf-8')
+setuptools.setup(
+    name='ithaca',
+    author='Ithaca team',
+    author_email='[email protected]',
+    version='0.1.0',
+    license='Apache License, Version 2.0',
+    description='Ithaca library for ancient text restoration and attribution.',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    packages=setuptools.find_packages(exclude=('train',)),
+    package_data={'': ['*.txt']},
+    install_requires=(here / 'requirements.txt').read_text().splitlines(),
+    extras_require={
+        'train': [
+            'optax',
+            'jaxline==0.0.5',
+            'tensorflow-datasets',
+        ]
+    },
+    classifiers=[
+        'Development Status :: 4 - Beta',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    ],
+)

train/README.md

@@ -0,0 +1,34 @@
+# Ithaca training code
+
+We recommend creating and activating a `conda` environment to ensure a clean
+environment where the correct package versions are installed below.
+```sh
+# Optional but recommended:
+conda create -n ithaca python==3.9
+conda activate ithaca
+```
+
+Clone this repository and enter its root directory. Install the full `ithaca`
+dependencies (including training), via:
+```sh
+git clone https://github.com/deepmind/ithaca
+cd ithaca
+pip install --editable .[train]
+cd train/
+```
+The `--editable` option links the `ithaca` installation to this repository, so
+that `import ithaca` will reflect any local modifications to the source code.
+
+Then, ensure you have TensorFlow installed. If you do not, install either the CPU or GPU version following the [instructions on the TensorFlow website](https://www.tensorflow.org/install/pip).
+While we use [Jax](https://github.com/google/jax) for training, TensorFlow is still needed for dataset loading.
+
+Next, ensure you have placed the dataset in `data/iphi.json`, note the wordlist and region mappings are also in that directory and may need to be replaced if they change in an updated version of the dataset. The dataset can be obtained from [I.PHI dataset](https://github.com/sommerschield/iphi).
+
+Finally, to run training, run:
+```sh
+./launch_local.sh
+```
+Alternatively, you can manually run:
+```sh
+python experiment.py --config=config.py --jaxline_mode=train --logtostderr
+```

train/config.py

@@ -0,0 +1,173 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Config for a Ithaca experiment."""
+
+from jaxline import base_config
+from ml_collections import config_dict
+
+
+def get_config():
+  """Return config object for training."""
+
+  config = base_config.get_base_config()
+
+  # Experiment config.
+  # Modify this to adapt to your custom distributed learning setup
+  local_batch_size = 1
+  num_devices = 1
+  config.train_batch_size = local_batch_size * num_devices
+
+  # Experiment config.
+  config.macros = config_dict.ConfigDict(
+      dict(
+          wordlist_size=35884,  # Keeping words with freq >10
+          context_char_max=768,
+          date_max=800,
+          date_min=-800,
+          date_interval=10,
+          date_bins=160,
+      ))
+  cm = config.macros  # Alias.
+
+  config.experiment_kwargs = config_dict.ConfigDict(
+      dict(
+          config=dict(
+              random_seed=4,
+              random_mode_train=config.get_ref('random_mode_train'),
+              random_mode_eval=config.get_ref('random_mode_eval'),
+              optimizer=dict(
+                  name='lamb',
+                  kwargs=dict(
+                      learning_rate=3e-4,
+                      weight_decay=0.,
+                      b2=0.999,
+                  ),
+                  # Set up the learning rate schedule.
+                  # factors='constant * linear_warmup * rsqrt_decay',
+                  warmup=4000,
+                  clip_adaptive=False,
+                  clip_level=0.,
+              ),
+              training=dict(
+                  batch_size=config.get_oneway_ref('train_batch_size')),
+              alphabet=dict(
+                  wordlist_path='data/iphi-wordlist.txt',
+                  wordlist_size=cm.get_ref('wordlist_size'),
+              ),
+              dataset=dict(
+                  dataset_path='data/iphi.json',
+                  region_main_path='data/iphi-region-main.txt',
+                  region_sub_path='data/iphi-region-sub.txt',
+                  context_char_min=50,
+                  context_char_max=cm.get_ref('context_char_max'),
+                  context_char_random=True,
+                  char_use_guess=True,
+                  char_mask_rate_min=0.,
+                  char_mask_rate_max=0.5,
+                  span_mask_eval_len=10,
+                  span_mask_ratio=0.15,
+                  span_mask_geometric_p=0.1,
+                  random_sentence_swap=0.25,
+                  random_word_delete=0.2,
+                  random_word_swap=0.,
+                  date_min=cm.get_ref('date_min'),
+                  date_max=cm.get_ref('date_max'),
+                  date_interval=cm.get_ref('date_interval'),
+                  date_bins=cm.get_ref('date_bins'),
+                  prepend_sos=1,
+                  repeat_train=-1,
+                  repeat_eval=10,
+                  black_list=[
+                      # 2334, 10, 293931, 14, 293752, 15, 293753, 16, 11,
+                      # 294468, 229647, 12, 291324, 291317, 17, 232697, 293754,
+                      # 1682, 1675, 1676, 1677, 1678, 1679, 1680, 1681, 291118,
+                      # 291320, 291319, 292366, 34, 291960, 35, 32, 346490, 27,
+                      # 292187, 291318, 19, 18, 37, 291321, 292189, 293756, 42,
+                      # 46, 232710, 39, 40, 41, 291322, 293757, 293327, 28,
+                      # 292194, 293326, 21, 293755, 291319, 291117, 38, 291959,
+                      # 31, 232705
+                  ],
+                  white_list=[]),
+              model=dict(
+                  word_char_emb_dim=256,
+                  emb_dim=512,
+                  mlp_dim=2048,
+                  num_layers=8,
+                  num_heads=4,
+                  vocab_char_size=34,
+                  vocab_word_size=cm.get_ref('wordlist_size') + 4,
+                  output_subregions=85,
+                  output_date=cm.get_ref('date_bins'),
+                  output_date_dist=True,
+                  region_date_pooling='first',
+                  use_output_mlp=True,
+                  max_len=cm.get_ref('context_char_max'),
+                  dropout_rate=0.1,
+                  attention_dropout_rate=0.1,
+                  use_bfloat16=False,
+                  model_type='bigbird',
+                  feature_combine_type='concat',
+                  posemb_combine_type='concat',
+              ),
+              loss=dict(
+                  date=dict(
+                      enabled=True,
+                      type='dist',
+                      weight_dist=1.25,
+                      weight_l1=0.,
+                      label_smoothing=0.,
+                      step_start=0,
+                      step_end=0,
+                  ),
+                  region=dict(
+                      enabled=True,
+                      weight=2.,
+                      label_smoothing=0.1,
+                      step_start=0,
+                      step_end=0,
+                  ),
+                  mask=dict(
+                      enabled=True,
+                      weight=3.,
+                      label_smoothing=0.05,
+                      step_start=0,
+                      step_end=0,
+                  ),
+                  nsp=dict(
+                      enabled=True,
+                      weight=0.01,
+                      step_start=0,
+                      step_end=0,
+                  )),
+              evaluation=dict(
+                  use_jit=True,
+                  batch_size=1,
+                  mode='valid',
+                  store_model_log=False,
+                  store_model_log_steps=100,
+              ),
+          ),))
+
+  # Training loop config.
+  config.training_steps = 1_000_000
+  config.log_train_data_interval = 10
+  config.save_checkpoint_interval = 300
+  config.best_model_eval_metric = 'score/eval'
+  config.checkpoint_dir = '/tmp/ithaca_checkpoints'
+  config.train_checkpoint_all_hosts = False
+
+  # Prevents accidentally setting keys that aren't recognized (e.g. in tests).
+  config.lock()
+
+  return config

train/data/README

@@ -0,0 +1,2 @@
+Put iphi.json in this directory before training.
+See train/README.md for details.

train/data/iphi-region-main.txt

@@ -0,0 +1,15 @@
+Asia Minor_1702;48510
+Attica (IG I-III)_1701;27215
+Aegean Islands, incl. Crete (IG XI-[XIII])_1699;25309
+Central Greece (IG VII-IX)_1698;14789
+Egypt, Nubia and Cyrenaïca_1695;12384
+Sicily, Italy, and the West (IG XIV)_1696;10108
+Greater Syria and the East_1693;9191
+Thrace and the Lower Danube (IG X)_1697;8456
+Northern Greece (IG X)_1692;7932
+Peloponnesos (IG IV-[VI])_1690;6393
+Cyprus ([IG XV])_1680;4056
+North Shore of the Black Sea_1683;3687
+North Africa_1614;287
+Upper Danube_1644;232
+Unknown Provenances_884;2

train/data/iphi-region-sub.txt

@@ -0,0 +1,88 @@
+Attica_1700;25595
+Egypt and Nubia_1694;11654
+Ionia_1688;10566
+Caria _1682;6214
+Megaris, Oropia, and Boiotia (IG VII) _1691;5576
+Macedonia _1485;5552
+Italy, incl. Magna Graecia_1689;5365
+Thrace and Moesia Inferior _1687;5257
+Rhodes and S. Dodecanese (IG XII,1) _1627;4957
+Cos and Calymna (IG XII,4)_1646;4487
+Phrygia_1679;4278
+Delphi _1272;4159
+unspecified subregion_1681;4056
+Sicily, Sardinia, and neighboring Islands_1686;4041
+unspecified subregion_1684;3687
+Lydia_1654;3621
+Thessaly (IG IX,2) _899;3586
+Syria and Phoenicia_1676;3444
+Delos (IG XI and ID) _1672;3431
+Mysia [Kaïkos], Pergamon_1674;3301
+Pisidia_1671;3110
+Arabia_1668;3064
+Scythia Minor _1675;2842
+Lycia_1669;2687
+Crete _474;2652
+Cilicia and Isauria_1662;2423
+Epeiros, Illyria, and Dalmatia _1463;2380
+Bithynia_1666;2253
+Samos (IG XII,6) _1665;1853
+Pamphylia_1651;1797
+Euboia (IG XII,9) _1653;1792
+Galatia_1663;1772
+Lycaonia_1661;1693
+Saronic Gulf, Corinthia, and the Argolid (IG IV) _1667;1620
+Lakonia and Messenia (IG V,1) _1658;1575
+Doric Sporades (IG XII,3) _1547;1529
+Cyclades, excl. Delos (IG XII,5) _1585;1526
+Mysia_1656;1475
+Phokis, Lokris, Aitolia, Akarnania, and Ionian Islands (IG IX,1) _1659;1449
+Pontus and Paphlagonia_1377;1369
+Palaestina_1657;1363
+Epidauria (IG IV²,1) _1643;1119
+Northern Aegean (IG XII,8) _1596;1099
+Elis _1647;1069
+Lesbos, Nesos, and Tenedos (IG XII,2) _1554;870
+Eleusis_1640;846
+Rhamnous_1639;774
+Mesopotamia_1626;740
+Cyrenaïca_1633;730
+Arkadia (IG V,2) _1632;702
+Troas_1631;583
+Amorgos and vicinity (IG XII,7) _1443;565
+Chios _1624;548
+Caria, Rhodian Peraia _1617;518
+Gallia_1635;451
+Aeolis_1012;396
+Dacia _1673;357
+Cappadocia_1607;344
+Achaia_1621;308
+Commagene_1590;229
+Africa Proconsularis_1589;189
+Hispania and Lusitania_1595;160
+Raetia, Noricum, and Pannonia _1574;116
+Moesia Superior _1641;116
+Bactria, Sogdiana_1558;104
+Babylonia_1535;77
+Tripolitania_1578;68
+Mysia [Upper Kaïkos] / Lydia_1482;50
+Susiana_1480;45
+Unknown Provenance_1477;43
+Germania_1502;42
+Armenia _1497;41
+Arabian Peninsula_1453;34
+unspecified subregion_1439;27
+Britannia_1445;22
+Doris_1434;19
+Iberia and Colchis_1421;19
+Arachosia, Drangiana_1417;17
+Persis_1399;15
+Mauretania Caesariensis_1381;14
+Media_1324;12
+Byzacena_1474;9
+Mauretania Tingitana_1090;4
+Numidia_1147;3
+unspecified subregion_885;2
+Hyrcania, Parthia_1146;2
+Osrhoene_634;1
+Carmania_881;1

train/dataloader.py

@@ -0,0 +1,409 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dataloader functions."""
+
+import json
+import random
+import re
+
+from absl import logging
+from ithaca.util.dates import date_range_to_dist
+from ithaca.util.text import random_mask_span
+from ithaca.util.text import random_sentence_swap
+from ithaca.util.text import random_word_delete
+from ithaca.util.text import random_word_swap
+from ithaca.util.text import text_to_idx
+from ithaca.util.text import text_to_word_idx
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+
+def generate_sample(config, alphabet, region_map, sample, mode='train'):
+  """Generates a new TF dataset sample."""
+
+  # Get text
+  text = sample['text']
+
+  # Next sentence prediction
+  sentences = text.split('.')
+  # Strip spaces
+  sentences = list(map(str.strip, sentences))
+  # Filter blank sentences
+  sentences = list(filter(None, sentences))
+  # Generate indexes
+  sentence_idx = np.arange(len(sentences), dtype=np.int32)
+
+  # Random sentence shuffling
+  if (mode == 'train' and config.random_sentence_swap > 0):
+    # Shuffle indexes
+    sentence_idx = random_sentence_swap(sentence_idx,
+                                        config.random_sentence_swap)
+    # Reshuffle sentences
+    sentences = np.array(sentences)[sentence_idx].tolist()
+
+  # Random word swap
+  if mode == 'train' and config.random_word_swap > 0:
+    sentences = [
+        random_word_swap(s, config.random_word_swap) for s in sentences
+    ]
+
+  # Random word delete
+  if mode == 'train' and config.random_word_delete > 0:
+    sentences = [
+        random_word_delete(s, config.random_word_delete) for s in sentences
+    ]
+
+  # Join text
+  text = '. '.join(sentences) + '.'
+
+  # Generate mask and labels
+  next_sentence_dots = np.array(
+      [pos for pos, char in enumerate(text[:-1]) if char == '.'],
+      dtype=np.int32)
+  next_sentence_mask = np.zeros(len(text), dtype=bool)
+  next_sentence_label = np.zeros(len(text), dtype=np.int32)
+  if sentence_idx.size > 1:
+    next_sentence_mask[next_sentence_dots] = True
+    next_sentence_label[next_sentence_dots] = (
+        sentence_idx[:-1] == (sentence_idx[1:] - 1))
+
+  # Computer start for prepending start of sentence character
+  start_sample_idx = int(config.prepend_sos)
+
+  if (mode in ['train', 'valid'] and config.context_char_random and
+      len(text) >= config.context_char_min):
+    # During training pick random context length
+    context_char_len = np.random.randint(
+        config.context_char_min,
+        min(len(text), config.context_char_max - start_sample_idx) + 1)
+
+    start_idx = 0
+    if context_char_len < len(text):
+      start_idx = np.random.randint(0, len(text) - context_char_len + 1)
+    text = text[start_idx:start_idx + context_char_len - start_sample_idx]
+    next_sentence_mask = next_sentence_mask[start_idx:start_idx +
+                                            context_char_len - start_sample_idx]
+    next_sentence_label = next_sentence_label[start_idx:start_idx +
+                                              context_char_len -
+                                              start_sample_idx]
+  elif (config.context_char_max and len(text) >
+        (config.context_char_max - start_sample_idx)):
+    # Clip text by maximum length
+    start_idx = np.random.randint(
+        0,
+        len(text) - (config.context_char_max - start_sample_idx) + 1)
+    text = text[start_idx:start_idx + config.context_char_max -
+                start_sample_idx]
+    next_sentence_mask = next_sentence_mask[start_idx:start_idx +
+                                            config.context_char_max -
+                                            start_sample_idx]
+    next_sentence_label = next_sentence_label[start_idx:start_idx +
+                                              config.context_char_max -
+                                              start_sample_idx]
+
+  # Prepend start of sentence character
+  if config.prepend_sos:
+    text = alphabet.sos + text
+    next_sentence_mask = [False] + next_sentence_mask
+    next_sentence_label = [0] + next_sentence_label
+
+  # Unmasked text
+  text_unmasked_idx = text_to_idx(text, alphabet)
+  text_unmasked_word_idx = text_to_word_idx(text, alphabet)
+
+  # Mask text
+  text_mask = np.zeros(len(text), dtype=bool)
+  if mode in ['train', 'valid']:
+    text_list = list(text)
+
+    # Non missing idx (avoid removing start of sentence character)
+    non_missing_idx = []
+    for i in range(start_sample_idx, len(text_list)):
+      if text_list[i] not in [alphabet.missing] + alphabet.punctuation:
+        non_missing_idx.append(i)
+
+    # Skip sample if there are no usable characters
+    if not non_missing_idx:
+      return
+
+    char_mask_idx = []
+    if config.char_mask_rate_max > 0.:
+      # Compute rate
+      char_mask_rate = np.random.uniform(config.char_mask_rate_min,
+                                         config.char_mask_rate_max)
+
+      # Fix masking in valid mode for comparing experiments
+      span_mask_geometric_p = config.span_mask_geometric_p
+      mask_num_total = int(char_mask_rate * len(non_missing_idx))
+      mask_num_span = int(mask_num_total * config.span_mask_ratio)
+      if mode == 'valid' and config.span_mask_eval_len > 0:
+        span_mask_geometric_p = None
+        mask_num_total = min(config.span_mask_eval_len, len(non_missing_idx))
+        mask_num_span = mask_num_total
+      mask_num_char = mask_num_total - mask_num_span
+
+      # Mask random indices
+      if mask_num_char > 0:
+        char_mask_idx = np.random.choice(
+            non_missing_idx, mask_num_char, replace=False).tolist()
+
+      # Mask random spans
+      if mask_num_span > 0:
+        count_span = 0
+        span_mask_idx = []
+        while (len(span_mask_idx) < mask_num_span and count_span < 10000):
+          span_mask_idx.extend(
+              random_mask_span(
+                  text,
+                  geometric_p=span_mask_geometric_p,
+                  limit_chars=mask_num_span - len(span_mask_idx)))
+          count_span += 1
+        char_mask_idx.extend(span_mask_idx)
+
+    # Mask text
+    for idx in set(char_mask_idx):
+      text_mask[idx] = True
+      text_list[idx] = alphabet.missing
+    text = ''.join(text_list)
+
+  # Text missing mask
+  text_missing_mask = np.array(list(text)) == alphabet.missing
+
+  # Convert to indices
+  text_idx = text_to_idx(text, alphabet)
+  text_idx_len = len(text_idx)
+  text_word_idx = text_to_word_idx(text, alphabet)
+  text_word_idx_len = len(text_word_idx)
+  assert text_idx_len == text_word_idx_len
+
+  # PHI id
+  phi_id = int(sample['id'])
+
+  # Map region ids to local ids
+  region_main_id = region_map['main']['ids_inv'][int(sample['region_main_id'])]
+  region_sub_id = region_map['sub']['ids_inv'][int(sample['region_sub_id'])]
+
+  # Dates
+  if (sample['date_min'] and sample['date_max'] and
+      int(sample['date_min']) <= int(sample['date_max']) and
+      int(sample['date_min']) >= config.date_min and
+      int(sample['date_max']) < config.date_max):
+    date_available = True
+    date_min = float(sample['date_min'])
+    date_max = float(sample['date_max'])
+    date_dist = date_range_to_dist(date_min, date_max, config.date_min,
+                                   config.date_max, config.date_interval,
+                                   config.date_bins)
+  else:
+    date_available = False
+    date_min = 0.
+    date_max = 0.
+    date_dist = date_range_to_dist(None, None, config.date_min, config.date_max,
+                                   config.date_interval, config.date_bins)
+
+  return {
+      'id': phi_id,  # 'text_str': text,
+      'text_char': text_idx,
+      'text_mask': text_mask,
+      'text_missing_mask': text_missing_mask,
+      'text_word': text_word_idx,
+      'text_len': text_idx_len,
+      'text_unmasked': text_unmasked_idx,
+      'text_unmasked_word': text_unmasked_word_idx,
+      'next_sentence_mask': next_sentence_mask,
+      'next_sentence_label': next_sentence_label,
+      'region_main_id': region_main_id,
+      'region_sub_id': region_sub_id,
+      'date_available': date_available,
+      'date_min': date_min,
+      'date_max': date_max,
+      'date_dist': date_dist,
+  }
+
+
+def loader_tf(batch_size,
+              config,
+              region_map,
+              alphabet=None,
+              dataset_file=None,
+              mode='train'):
+  """TF dataloader."""
+  # Load dataset
+  dataset_tmp = {int(d['id']): d for d in json.load(dataset_file)}
+  logging.info('Loaded dataset inscriptions: %d.', len(dataset_tmp))
+
+  # Check if white_list enabled
+  if hasattr(config, 'white_list') and config.white_list:
+    dataset = []
+    for d in dataset_tmp.values():
+      if int(d['id']) in config.white_list:
+        dataset.append(d)
+    del dataset_tmp
+  else:
+    # Find duplicate inscriptions
+    rev_dataset = {}
+    black_list = set()
+    if hasattr(config, 'black_list') and config.black_list:
+      logging.info('Ignore list inscriptions: %d.', len(config.black_list))
+      black_list.update(config.black_list)
+
+    for key in sorted(dataset_tmp.keys()):
+      value = dataset_tmp[key]
+      rev_dataset.setdefault(value['text'], set()).add(key)
+      if len(rev_dataset[value['text']]) > 1:
+        black_list.add(int(value['id']))
+    del rev_dataset
+    logging.info('Inscriptions filtered: %d.', len(black_list))
+
+    # Create deduplicated dataset
+    dataset = []
+    for d in dataset_tmp.values():
+      if int(d['id']) not in black_list:
+        dataset.append(d)
+    del dataset_tmp
+    del black_list
+
+  logging.info('Final dataset inscriptions: %d.', len(dataset))
+
+  # Breaks dataset correlated order.
+  random.shuffle(dataset)
+
+  # Sample generator function
+  def generate_samples():
+
+    dataset_idxs = list(range(len(dataset)))
+    random.shuffle(dataset_idxs)
+    for dataset_i in dataset_idxs:
+      sample = dataset[dataset_i]
+
+      # Skip if region does not exist in map
+      if (int(sample['region_main_id']) not in region_map['main']['ids_inv'] or
+          int(sample['region_sub_id']) not in region_map['sub']['ids_inv']):
+        continue
+
+      # Replace guess signs with missing chars
+      if hasattr(config, 'char_use_guess') and not config.char_use_guess:
+        sample['text'] = re.sub(r'\[(.*?)\]', lambda m: '-' * len(m.group(1)),
+                                sample['text'])
+      sample['text'] = sample['text'].replace(alphabet.sog,
+                                              '').replace(alphabet.eog, '')
+
+      # Filter by text length
+      if len(sample['text'].replace(alphabet.missing,
+                                    '')) < config.context_char_min:
+        continue
+
+      # Last digit 3 -> test, 4 -> valid, the rest are the training set
+      sample_id = int(sample['id'])
+      if ((sample_id % 10 == 3 and mode == 'test') or
+          (sample_id % 10 == 4 and mode == 'valid') or
+          (sample_id % 10 != 3 and sample_id % 10 != 4 and mode == 'train') or
+          (hasattr(config, 'white_list') and config.white_list)):
+        s = generate_sample(config, alphabet, region_map, sample, mode=mode)
+        if s:
+          yield s
+
+  # Create dataset from generator.
+  with tf.device('/cpu:0'):
+    ds = tf.data.Dataset.from_generator(
+        generate_samples,
+        output_signature={
+            'id':
+                tf.TensorSpec(shape=(), dtype=tf.int32),
+            'text_char':
+                tf.TensorSpec(shape=(None), dtype=tf.int32),
+            'text_mask':
+                tf.TensorSpec(shape=(None), dtype=tf.bool),
+            'text_missing_mask':
+                tf.TensorSpec(shape=(None), dtype=tf.bool),
+            'text_word':
+                tf.TensorSpec(shape=(None), dtype=tf.int32),
+            'text_unmasked':
+                tf.TensorSpec(shape=(None), dtype=tf.int32),
+            'text_unmasked_word':
+                tf.TensorSpec(shape=(None), dtype=tf.int32),
+            'next_sentence_mask':
+                tf.TensorSpec(shape=(None), dtype=tf.bool),
+            'next_sentence_label':
+                tf.TensorSpec(shape=(None), dtype=tf.int32),
+            'text_len':
+                tf.TensorSpec(shape=(), dtype=tf.int32),
+            'region_main_id':
+                tf.TensorSpec(shape=(), dtype=tf.int32),
+            'region_sub_id':
+                tf.TensorSpec(shape=(), dtype=tf.int32),
+            'date_available':
+                tf.TensorSpec(shape=(), dtype=tf.bool),
+            'date_min':
+                tf.TensorSpec(shape=(), dtype=tf.float32),
+            'date_max':
+                tf.TensorSpec(shape=(), dtype=tf.float32),
+            'date_dist':
+                tf.TensorSpec(shape=(config.date_bins), dtype=tf.float32),
+        })
+
+  # Shuffle and repeat.
+  if mode == 'train':
+    if config.repeat_train == -1:
+      ds = ds.repeat()
+    elif config.repeat_train >= 1:
+      ds = ds.repeat(config.repeat_train)
+  else:
+    if config.repeat_eval == -1:
+      ds = ds.repeat()
+    elif config.repeat_eval >= 1:
+      ds = ds.repeat(config.repeat_eval)
+
+  # Batch and pad.
+  max_len = config.context_char_max
+  ds = ds.padded_batch(
+      batch_size,
+      padded_shapes={
+          'id': [],
+          'text_char': [max_len],
+          'text_mask': [max_len],
+          'text_missing_mask': [max_len],
+          'text_word': [max_len],
+          'text_unmasked': [max_len],
+          'text_unmasked_word': [max_len],
+          'next_sentence_mask': [max_len],
+          'next_sentence_label': [max_len],
+          'text_len': [],
+          'region_main_id': [],
+          'region_sub_id': [],
+          'date_available': [],
+          'date_min': [],
+          'date_max': [],
+          'date_dist': [config.date_bins]
+      },
+      padding_values={
+          'id': 0,
+          'text_char': alphabet.pad_idx,
+          'text_mask': False,
+          'text_missing_mask': True,
+          'text_word': alphabet.pad_idx,
+          'text_unmasked': alphabet.pad_idx,
+          'text_unmasked_word': alphabet.pad_idx,
+          'next_sentence_mask': False,
+          'next_sentence_label': 0,
+          'text_len': 0,
+          'region_main_id': 0,
+          'region_sub_id': 0,
+          'date_available': False,
+          'date_min': 0.,
+          'date_max': 0.,
+          'date_dist': 0.
+      })
+
+  return ds

train/experiment.py

@@ -0,0 +1,679 @@
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Ithaca: Restoring and attributing ancient texts with deep neural networks."""
+
+import bz2
+import distutils
+import functools
+import glob
+import os
+import pickle
+
+from absl import app
+from absl import flags
+from absl import logging
+import dataloader
+from ithaca.models.model import Model
+from ithaca.util.alphabet import GreekAlphabet
+from ithaca.util.loss import categorical_kl_divergence
+from ithaca.util.loss import cross_entropy_label_smoothing_loss
+from ithaca.util.loss import cross_entropy_loss
+from ithaca.util.loss import cross_entropy_mask_loss
+from ithaca.util.loss import date_loss_l1
+from ithaca.util.optim import adaptive_grad_clip
+from ithaca.util.optim import linear_warmup_and_sqrt_decay
+from ithaca.util.optim import linear_weight
+from ithaca.util.region_names import load_region_maps
+import jax
+import jax.numpy as jnp
+from jaxline import experiment
+from jaxline import platform
+from jaxline import utils as jl_utils
+import numpy as np
+import optax
+import tensorflow_datasets.public_api as tfds
+
+FLAGS = flags.FLAGS
+
+
+class Experiment(experiment.AbstractExperiment):
+  """Ithaca experiment."""
+
+  # Holds a map from object properties that will be checkpointed to their name
+  # within a checkpoint. Currently it is assume that these are all sharded
+  # device arrays.
+  CHECKPOINT_ATTRS = {
+      '_params': 'params',
+      '_opt_state': 'opt_state',
+  }
+
+  def __init__(self, mode, init_rng, config):
+    """Initializes experiment."""
+
+    super(Experiment, self).__init__(mode=mode)
+    self.mode = mode
+    self.init_rng = init_rng
+    self.config = config
+
+    # Same random key on each device.
+    self._rng_key = jl_utils.bcast_local_devices(self.init_rng)
+
+    # Checkpointed experiment state.
+    self._params = None
+    self._opt_state = None
+
+    # Input pipelines.
+    self._train_input = None
+    self._eval_input = None
+
+    # Forward and update functions.
+    self.forward = Model(**self.config.model)
+    self._update_func = jax.pmap(
+        self._update_func, axis_name='i', donate_argnums=(0, 1))
+
+    self._learning_rate_fn = functools.partial(
+        linear_warmup_and_sqrt_decay,
+        max_lr=self.config.optimizer.kwargs.learning_rate,
+        warmup_steps=self.config.optimizer.warmup)
+
+    self._opt_init, self._opt_update = self.optimizer()
+
+    if 'use_jit' in self.config.evaluation and self.config.evaluation.use_jit:
+      self._eval_batch = jax.jit(self._eval_batch)
+
+    # Create alphabet
+    alphabet_kwargs = dict(self.config.alphabet)
+    wordlist_path = alphabet_kwargs.pop('wordlist_path')
+    with open(wordlist_path, 'r') as f:
+      self._alphabet = GreekAlphabet(wordlist_file=f, **alphabet_kwargs)
+
+    # Create region mapping
+    self._region_map = {'main': None, 'sub': None}
+    if self.config.dataset.region_main_path:
+      with open(self.config.dataset.region_main_path, 'r') as f:
+        self._region_map['main'] = load_region_maps(f)
+    if self.config.dataset.region_sub_path:
+      with open(self.config.dataset.region_sub_path, 'r') as f:
+        self._region_map['sub'] = load_region_maps(f)
+
+  def optimizer(self):
+    config_opt = self.config.optimizer
+
+    kwargs = config_opt.kwargs.to_dict()
+    kwargs['learning_rate'] = self._learning_rate_fn
+    opt = getattr(optax, config_opt.name)(**kwargs)
+
+    if hasattr(config_opt, 'clip_adaptive') and config_opt.clip_adaptive:
+      if config_opt.clip_level > 0.:
+        opt = optax.chain(adaptive_grad_clip(config_opt.clip_level), opt)
+    elif config_opt.clip_level > 0.:
+      opt = optax.chain(optax.clip_by_global_norm(config_opt.clip_level), opt)
+    return opt
+
+  #  _             _
+  # | |_ _ __ __ _(_)_ __
+  # | __| '__/ _` | | '_ \
+  # | |_| | | (_| | | | | |
+  #  \__|_|  \__,_|_|_| |_|
+  #
+
+  def step(self, global_step, rng, **unused_args):
+    """See base class."""
+
+    if self._train_input is None:
+      self._initialize_train(rng)
+
+    batch = next(self._train_input)
+    (self._params, self._opt_state, scalars) = (
+        self._update_func(self._params, self._opt_state, global_step, batch,
+                          rng))
+
+    scalars = jl_utils.get_first(scalars)
+    return scalars
+
+  def _initialize_train(self, rng):
+    # Check we haven't already restored params
+    if self._params is None:
+      logging.info(
+          'Initializing parameters rather than restoring from checkpoint.')
+      batch = next(self._build_train_input())
+
+      rng = jl_utils.get_first(rng)
+      params_rng, dropout_rng = jax.random.split(rng)
+      params_rng = jl_utils.bcast_local_devices(params_rng)
+      dropout_rng = jl_utils.bcast_local_devices(dropout_rng)
+      init_net = jax.pmap(
+          functools.partial(self.forward.init, is_training=True))
+      self._params = init_net({
+          'params': params_rng,
+          'dropout': dropout_rng
+      },
+                              text_char=batch['text_char'],
+                              text_word=batch['text_word'])
+
+      init_opt = jax.pmap(self._opt_init)
+      self._opt_state = init_opt(self._params)
+
+      self._train_input = jl_utils.py_prefetch(self._build_train_input)
+      self._train_input = jl_utils.double_buffer_on_gpu(self._train_input)
+
+  def _build_train_input(self):
+    """See base class."""
+    num_devices = jax.device_count()
+    global_batch_size = self.config.training.batch_size
+    per_device_batch_size, ragged = divmod(global_batch_size, num_devices)
+    logging.info(
+        'num_devices: %d, per_device_batch_size: %d, global_batch_size: %d',
+        num_devices, per_device_batch_size, global_batch_size)
+
+    if ragged:
+      raise ValueError(
+          f'Global batch size {global_batch_size} must be divisible by '
+          f'num devices {num_devices}')
+
+    config_dataset = self.config.dataset
+    with open(config_dataset.dataset_path) as dataset_file:
+      ds = dataloader.loader_tf(
+          per_device_batch_size,
+          config_dataset,
+          self._region_map,
+          alphabet=self._alphabet,
+          dataset_file=dataset_file,
+          mode='train')
+
+    ds = ds.batch(jax.local_device_count())
+    return iter(tfds.as_numpy(ds))
+
+  def _loss_fn(self, params, batch, global_step, rng):
+    text_char = batch['text_char']
+    text_word = batch['text_word']
+    text_unmasked = batch['text_unmasked']
+    text_mask = batch['text_mask']
+    next_sentence_mask = batch['next_sentence_mask']
+    next_sentence_label = batch['next_sentence_label']
+    subregion = batch['region_sub_id']
+    date_min = batch['date_min']
+    date_max = batch['date_max']
+    date_dist = batch['date_dist']
+    date_available = batch['date_available']
+    eps = 1e-6
+
+    (date_pred, subregion_logits, mask_logits, nsp_logits) = self.forward.apply(
+        params,
+        text_char=text_char,
+        text_word=text_word,
+        text_char_onehot=None,
+        text_word_onehot=None,
+        is_training=True,
+        rngs={'dropout': rng})
+
+    date_loss = 0.
+    subregion_loss = 0.
+    subregion_accuracy = 0.
+    mask_loss = 0.
+    mask_accuracy = 0.
+    nsp_loss = 0.
+    nsp_accuracy = 0.
+
+    # Date loss
+    if self.config.loss.date.enabled:
+      if self.config.loss.date.label_smoothing > 0:
+        date_dist_prob = jnp.exp(date_dist)  # logprob to prob
+        date_dist_prob_smooth = date_dist_prob * jax.random.uniform(
+            rng,
+            shape=date_dist_prob.shape,
+            dtype=date_dist_prob.dtype,
+            minval=1 - self.config.loss.date.label_smoothing,
+            maxval=1 + self.config.loss.date.label_smoothing)
+        date_dist_prob_smooth /= date_dist_prob_smooth.sum(axis=-1)[:,
+                                                                    jnp.newaxis]
+
+        date_dist_prob_smooth = jnp.clip(date_dist_prob_smooth, 1e-6, 1)
+        date_dist = jnp.log(date_dist_prob_smooth)
+
+      date_loss = 0.
+      if 'l1' in self.config.loss.date.type.split('+'):
+        date_pred_x = jnp.arange(
+            self.config.dataset.date_min +
+            self.config.dataset.date_interval / 2,
+            self.config.dataset.date_max +
+            self.config.dataset.date_interval / 2,
+            self.config.dataset.date_interval).reshape(-1, 1)
+        date_pred_val = jnp.dot(jax.nn.softmax(date_pred, axis=-1), date_pred_x)
+        date_loss_l1_ = jax.vmap(date_loss_l1)(date_pred_val, date_min,
+                                               date_max, date_available)
+        jnp.nan_to_num(date_loss_l1_, copy=False)
+        date_loss += (
+            jnp.mean(date_loss_l1_, axis=0) * self.config.loss.date.weight_l1)
+
+      if 'dist' in self.config.loss.date.type.split('+'):
+        date_loss_dist_ = categorical_kl_divergence(date_dist, date_pred)
+        date_loss_dist_ *= date_available
+        jnp.nan_to_num(date_loss_dist_, copy=False)
+        date_loss += (
+            jnp.mean(date_loss_dist_, axis=0) *
+            self.config.loss.date.weight_dist)
+
+      date_loss *= linear_weight(global_step, self.config.loss.date.step_start,
+                                 self.config.loss.date.step_end)
+
+    # Region and subregion loss
+    if self.config.loss.region.enabled:
+      subregion_loss = jnp.mean(
+          cross_entropy_label_smoothing_loss(
+              subregion_logits,
+              subregion,
+              label_smoothing=self.config.loss.region.label_smoothing), 0)
+      jnp.nan_to_num(subregion_loss, copy=False)
+      subregion_loss *= self.config.loss.region.weight
+      subregion_accuracy = jnp.mean(
+          jnp.argmax(subregion_logits, -1) == subregion)
+
+      w = linear_weight(global_step, self.config.loss.region.step_start,
+                        self.config.loss.region.step_end)
+      subregion_loss *= w
+
+    # Mask loss
+    if self.config.loss.mask.enabled:
+      mask_loss = jnp.sum(
+          cross_entropy_label_smoothing_loss(
+              mask_logits,
+              text_unmasked,
+              text_mask,
+              label_smoothing=self.config.loss.mask.label_smoothing), 1)  # [B]
+      assert mask_loss.ndim == 1
+      jnp.nan_to_num(mask_loss, copy=False)
+      mask_loss = jnp.mean(mask_loss, 0) * self.config.loss.mask.weight  # []
+      mask_all_accuracy = (jnp.argmax(mask_logits, -1) == text_unmasked).astype(
+          mask_logits.dtype)
+      mask_accuracy = jnp.divide(
+          jnp.sum(
+              jnp.multiply(mask_all_accuracy,
+                           text_mask.astype(mask_logits.dtype))),
+          jnp.sum(text_mask) + eps)
+
+      mask_loss *= linear_weight(global_step, self.config.loss.mask.step_start,
+                                 self.config.loss.mask.step_end)
+
+    # NSP loss
+    if self.config.loss.nsp.enabled:
+      nsp_loss = jnp.sum(
+          jax.vmap(jax.vmap(cross_entropy_mask_loss))(nsp_logits,
+                                                      next_sentence_label,
+                                                      next_sentence_mask),
+          1)  # [B]
+      assert nsp_loss.ndim == 1
+      jnp.nan_to_num(nsp_loss, copy=False)
+      nsp_loss = jnp.mean(nsp_loss, 0) * self.config.loss.nsp.weight
+      nsp_all_accuracy = (jnp.argmax(
+          nsp_logits, -1) == next_sentence_label).astype(nsp_logits.dtype)
+      nsp_accuracy = jnp.divide(
+          jnp.sum(
+              jnp.multiply(nsp_all_accuracy,
+                           next_sentence_mask.astype(nsp_logits.dtype))),
+          jnp.sum(next_sentence_mask) + eps)
+      nsp_loss *= linear_weight(global_step, self.config.loss.nsp.step_start,
+                                self.config.loss.nsp.step_end)
+
+    loss = date_loss + subregion_loss + mask_loss + nsp_loss
+    scaled_loss = loss / jax.device_count()
+    # NOTE: We use scaled_loss for grads and unscaled for logging.
+    return scaled_loss, (loss, date_loss, subregion_loss, subregion_accuracy,
+                         mask_loss, mask_accuracy, nsp_loss, nsp_accuracy)
+
+  def _update_func(self, params, opt_state, global_step, batch, rng):
+    """Applies an update to parameters and returns new state."""
+    # This function computes the gradient of the first output of loss_fn and
+    # passes through the other arguments unchanged.
+    grad_loss_fn = jax.grad(self._loss_fn, has_aux=True)
+    scaled_grads, (loss, date_loss, subregion_loss, subregion_accuracy,
+                   mask_loss, mask_accuracy, nsp_loss,
+                   nsp_accuracy) = grad_loss_fn(params, batch, global_step, rng)
+
+    scaled_grads = jax.tree_map(jnp.nan_to_num, scaled_grads)
+    grads = jl_utils.tree_psum(scaled_grads, axis_name='i')
+
+    # Compute and apply updates via our optimizer.
+    learning_rate = self._learning_rate_fn(global_step)
+    updates, opt_state = self._opt_update(grads, opt_state, params=params)
+    params = optax.apply_updates(params, updates)
+
+    # Scalars to log (note: we log the mean across all hosts/devices).
+    scalars = {
+        'loss/train': loss,
+        'loss/date': date_loss,
+        'loss/subregion': subregion_loss,
+        'loss/mask': mask_loss,
+        'loss/nsp': nsp_loss,
+        'accuracy/subregion': subregion_accuracy,
+        'accuracy/mask': mask_accuracy,
+        'accuracy/nsp': nsp_accuracy,
+        'opt/learning_rate': learning_rate,
+        'opt/grad_norm': optax.global_norm(grads),
+        'opt/param_norm': optax.global_norm(params),
+    }
+    scalars = jax.lax.pmean(scalars, axis_name='i')
+
+    return params, opt_state, scalars
+
+  #                  _
+  #   _____   ____ _| |
+  #  / _ \ \ / / _` | |
+  # |  __/\ V / (_| | |
+  #  \___| \_/ \__,_|_|
+  #
+
+  def evaluate(self, global_step, rng, **unused_kwargs):
+    """See base class."""
+
+    if self._eval_input is None:
+      self._initialize_eval()
+
+    global_step = np.array(jl_utils.get_first(global_step))
+    summary, outputs = self._eval_epoch(jl_utils.get_first(rng))
+
+    for k, v in summary.items():
+      summary[k] = np.array(v)
+
+    score = summary['score/eval']
+    logging.info('[Step %d] eval_score=%.2f', global_step, score)
+
+    # Log outputs
+    checkpoint_dir = jl_utils.get_checkpoint_dir(FLAGS.config,
+                                                 jax.process_index())
+    outputs_path = os.path.join(checkpoint_dir, 'best_outputs.pkl.bz2')
+    score_path = os.path.join(checkpoint_dir, 'best_score.txt')
+    model_log_path = os.path.join(checkpoint_dir, 'model_log')
+    best_model_log_path = os.path.join(checkpoint_dir, 'best_model_log')
+
+    # Check for preexisting outputs
+    best_score = None
+    best_step = None
+    if os.path.exists(score_path):
+      with open(score_path, 'r') as f:
+        tok = f.read().strip().split(' ')
+        best_step = int(tok[0])
+        best_score = float(tok[1])
+
+    # Store outputs if score is better
+    if best_score is None or (score > best_score and global_step > best_step):
+      best_score = score
+
+      with open(score_path, 'w') as f:
+        f.write(f'{global_step} {best_score}')
+
+      with open(outputs_path, 'wb') as f:
+        outputs_pkl = pickle.dumps(outputs, protocol=2)
+        outputs_pkl_bz2 = bz2.compress(outputs_pkl)
+        f.write(outputs_pkl_bz2)
+
+      if self.config.evaluation.store_model_log:
+        if os.path.isdir(best_model_log_path):
+          map(os.remove, glob.glob(best_model_log_path + '/*'))
+        else:
+          os.makedirs(best_model_log_path)
+        distutils.dir_util.copy_tree(model_log_path, best_model_log_path)
+
+      logging.info('[Step %d] Writing eval outputs: %s.', global_step,
+                   outputs_path)
+
+    # Log best score
+    summary['score/eval_best'] = best_score
+
+    return summary
+
+  def _initialize_eval(self):
+    self._eval_input = jl_utils.py_prefetch(self._build_eval_input)
+
+  def _build_eval_input(self):
+    """Builds the evaluation input pipeline."""
+    config_dataset = self.config.dataset
+    with open(config_dataset.dataset_path) as dataset_file:
+      ds = dataloader.loader_tf(
+          self.config.evaluation.batch_size,
+          config_dataset,
+          self._region_map,
+          alphabet=self._alphabet,
+          dataset_file=dataset_file,
+          mode=self.config.evaluation.mode)
+
+    return iter(tfds.as_numpy(ds))
+
+  def _eval_batch(self, params, batch, rng):
+    """Evaluates a batch."""
+    phi_id = batch['id']
+    text_char = batch['text_char']
+    text_word = batch['text_word']
+    text_unmasked = batch['text_unmasked']
+    text_mask = batch['text_mask']
+    next_sentence_mask = batch['next_sentence_mask']
+    next_sentence_label = batch['next_sentence_label']
+    subregion = batch['region_sub_id']
+    date_min = batch['date_min']
+    date_max = batch['date_max']
+    date_dist = batch['date_dist']
+    date_available = batch['date_available']
+
+    # with hlogging.context() as log:
+    (date_pred, subregion_logits, mask_logits, nsp_logits) = self.forward.apply(
+        params,
+        text_char=text_char,
+        text_word=text_word,
+        text_char_onehot=None,
+        text_word_onehot=None,
+        is_training=False,
+        rngs={'dropout': rng})
+
+    # Log model weights
+    model_log = {}
+
+    subregion_loss = 0.
+    subregion_accuracy = 0.
+    date_loss = 0.
+    date_l1_loss = 0.
+    nsp_loss = 0.
+    nsp_accuracy = 0.
+    # eps = 1e-6
+
+    date_count = 0
+    mask_count = 0
+    nsp_count = 0
+
+    # Date loss
+    if self.config.loss.date.enabled:
+      date_pred_x = jnp.arange(
+          self.config.dataset.date_min + self.config.dataset.date_interval / 2,
+          self.config.dataset.date_max + self.config.dataset.date_interval / 2,
+          self.config.dataset.date_interval).reshape(-1, 1)
+      date_pred_val = jnp.dot(jax.nn.softmax(date_pred, axis=-1), date_pred_x)
+      date_l1_loss = jnp.sum(
+          jax.vmap(date_loss_l1)(date_pred_val, date_min, date_max,
+                                 date_available),
+          axis=0)
+
+      if 'l1' in self.config.loss.date.type.split('+'):
+        date_loss += date_l1_loss * self.config.loss.date.weight_l1
+
+      if 'dist' in self.config.loss.date.type.split('+'):
+        date_loss_dist_ = categorical_kl_divergence(date_dist, date_pred)
+        date_loss_dist_ *= date_available
+        date_loss += (
+            jnp.sum(date_loss_dist_, axis=0) *
+            self.config.loss.date.weight_dist)
+
+      date_count = jnp.sum(date_available)
+
+    # Region and subregion loss
+    if self.config.loss.region.enabled:
+      subregion_loss = jnp.sum(
+          cross_entropy_loss(subregion_logits, subregion), 0)
+      subregion_loss *= self.config.loss.region.weight
+      subregion_accuracy = jnp.mean(
+          jnp.argmax(subregion_logits, -1) == subregion)
+
+    # Mask loss
+    if self.config.loss.mask.enabled:
+      mask_loss = jnp.sum(
+          cross_entropy_label_smoothing_loss(
+              mask_logits, text_unmasked, text_mask, label_smoothing=0),
+          1)  # [B]
+      # mask_loss /= jnp.sum(text_mask, axis=1) + eps  # [B]
+      assert mask_loss.ndim == 1
+      mask_loss = jnp.mean(mask_loss, 0) * self.config.loss.mask.weight  # []
+
+      mask_all_accuracy = (jnp.argmax(mask_logits, -1) == text_unmasked).astype(
+          mask_logits.dtype)
+      mask_accuracy = jnp.sum(
+          jnp.multiply(mask_all_accuracy, text_mask.astype(mask_logits.dtype)))
+      mask_count = jnp.sum(text_mask)
+
+    # NSP loss
+    if self.config.loss.nsp.enabled:
+      nsp_loss = jnp.sum(
+          jax.vmap(jax.vmap(cross_entropy_mask_loss))(nsp_logits,
+                                                      next_sentence_label,
+                                                      next_sentence_mask),
+          1)  # [B]
+      assert nsp_loss.ndim == 1
+      nsp_loss = jnp.sum(nsp_loss, 0) * self.config.loss.nsp.weight
+      nsp_all_accuracy = (jnp.argmax(
+          nsp_logits, -1) == next_sentence_label).astype(nsp_logits.dtype)
+      nsp_accuracy = jnp.sum(
+          jnp.multiply(nsp_all_accuracy,
+                       next_sentence_mask.astype(nsp_logits.dtype)))
+      nsp_count = jnp.sum(next_sentence_mask)
+
+    # Outputs
+    scalars = {
+        'score/eval':
+            (mask_accuracy + subregion_accuracy - date_l1_loss * 0.01),
+        'loss/eval': mask_loss + date_loss + subregion_loss,
+        'loss/date': date_loss,
+        'loss/date_l1': date_l1_loss,
+        'loss/subregion': subregion_loss,
+        'loss/mask': mask_loss,
+        'loss/nsp': nsp_loss,
+        'count/date': date_count,
+        'count/nsp': nsp_count,
+        'count/mask': mask_count,
+        'accuracy/subregion': subregion_accuracy,
+        'accuracy/mask': mask_accuracy,
+        'accuracy/nsp': nsp_accuracy,
+    }
+
+    outputs = {
+        'outputs/id': phi_id,
+        'outputs/date_pred': date_pred.astype('float16'),
+        'outputs/date_min': date_min,
+        'outputs/date_max': date_max,
+        'outputs/date_dist': date_dist.astype('float16'),
+        'outputs/date_available': date_available,
+        'outputs/subregion_logits': subregion_logits.astype('float16'),
+        'outputs/subregion': subregion,
+    }
+
+    return scalars, outputs, model_log
+
+  def _eval_epoch(self, rng):
+    """Evaluates an epoch."""
+    summary = {}
+    outputs = {}
+    total_num_sequences = 0
+
+    # Prepare directories for storing model log
+    checkpoint_dir = jl_utils.get_checkpoint_dir(FLAGS.config,
+                                                 jax.process_index())
+    model_log_path = os.path.join(checkpoint_dir, 'model_log')
+    if self.config.evaluation.store_model_log:
+      if os.path.isdir(model_log_path):
+        map(os.remove, glob.glob(model_log_path + '/*'))
+      else:
+        os.makedirs(model_log_path)
+
+    # Checkpoints broadcast for each local device
+    params = jl_utils.get_first(self._params)
+
+    # Model log buffer initialisation
+    model_log_buffer = []
+
+    def _flush_model_log_buffer(model_log_buffer):
+      """Writes model log to bz2 pickle files."""
+      while model_log_buffer:
+        model_log_batch_path, model_log_pkl_bz2 = model_log_buffer.pop(0)
+        with open(model_log_batch_path, 'wb') as f:
+          f.write(model_log_pkl_bz2)
+
+    # Converting to numpy here allows us to reset the generator
+    for batch in self._eval_input():
+      # Make sure that the input has batch_dim=1
+      assert batch['text_char'].shape[0] == 1
+
+      summary_batch, outputs_batch, model_log_batch = self._eval_batch(
+          params, batch, rng)
+
+      # Append batch values to dictionary
+      for k, v in summary_batch.items():
+        summary[k] = summary.get(k, 0) + v
+      for k, v in outputs_batch.items():
+        outputs.setdefault(k, []).append(v)
+
+      total_num_sequences += self.config.evaluation.batch_size
+
+      # Store model log per batch
+      if self.config.evaluation.store_model_log:
+        # Append to buffer
+        model_log_batch_path = os.path.join(
+            model_log_path,
+            str(outputs_batch['outputs/id'][0]) + '.pkl.bz2')
+        model_log_pkl = pickle.dumps(model_log_batch, protocol=2)
+        model_log_pkl_bz2 = bz2.compress(model_log_pkl)
+        model_log_buffer += [(model_log_batch_path, model_log_pkl_bz2)]
+
+        # Flush model log buffer
+        if (len(model_log_buffer) %
+            self.config.evaluation.store_model_log_steps == 0):
+          _flush_model_log_buffer(model_log_buffer)
+
+    # Flush remaining model log buffer
+    if self.config.evaluation.store_model_log:
+      _flush_model_log_buffer(model_log_buffer)
+
+    # Normalise and concatenate
+    summary['loss/date'] /= summary['count/date']
+    summary['loss/date_l1'] /= summary['count/date']
+
+    summary['loss/mask'] /= summary['count/mask']
+    summary['accuracy/mask'] /= summary['count/mask']
+
+    summary['loss/nsp'] /= summary['count/nsp']
+    summary['accuracy/nsp'] /= summary['count/nsp']
+
+    summary['loss/subregion'] /= total_num_sequences
+    summary['accuracy/subregion'] /= total_num_sequences
+
+    summary['score/eval'] = (
+        summary['accuracy/mask'] + summary['accuracy/subregion'] -
+        summary['loss/date_l1'] * 0.01)
+    summary['loss/eval'] = (
+        summary['loss/mask'] + summary['loss/date'] + summary['loss/subregion'])
+
+    for k, v in outputs.items():
+      outputs[k] = np.concatenate(v, axis=0)
+
+    return summary, outputs
+
+
+if __name__ == '__main__':
+  flags.mark_flag_as_required('config')
+  app.run(functools.partial(platform.main, Experiment))

train/launch_local.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+# Copyright 2021 the Ithaca Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Command for running training script. Run from within train/ directory
+# after first installing the ithaca package.
+#
+# See README.md in this train/ directory for details.
+
+python experiment.py --config=config.py --jaxline_mode=train --logtostderr