# coding=utf-8
# Copyright 2019 The Google Research 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.

# Lint as: python3
"""Utility functions for running inference with a LaserTagger model."""

from __future__ import absolute_import
from __future__ import division

from __future__ import print_function

from typing import Mapping, Sequence, Text

import bert_example
import tagging


class LaserTaggerPredictor(object):
  """Class for computing and realizing predictions with LaserTagger."""

  def __init__(self, tf_predictor,
               example_builder,
               label_map):
    """Initializes an instance of LaserTaggerPredictor.

    Args:
      tf_predictor: Loaded Tensorflow model.
      example_builder: BERT example builder.
      label_map: Mapping from tags to tag IDs.
    """
    self._predictor = tf_predictor
    self._example_builder = example_builder
    self._id_2_tag = {
        tag_id: tagging.Tag(tag) for tag, tag_id in label_map.items()
    }

  def predict(self, sources):
    """Returns realized prediction for given sources."""
    example = self._example_builder.build_bert_example(sources)
    if example is None:
      raise ValueError("Example couldn't be built.")

    # Predict tag IDs.
    keys = ['input_ids', 'input_mask', 'segment_ids']
    out = self._predictor({key: [example.features[key]] for key in keys})
    predicted_ids = out['pred'][0].tolist()
    # Realize output.
    example.features['labels'] = predicted_ids
    # Mask out the begin and the end token.
    example.features['labels_mask'] = [0] + [1] * (len(predicted_ids) - 2) + [0]
    labels = [
        self._id_2_tag[label_id] for label_id in example.get_token_labels()
    ]
    return example.editing_task.realize_output(labels)

@@ -1562,24 +1562,140 @@
 sources.%22%22%22%0A
+    # Predict tag IDs.%0A    keys = %5B'input_ids', 'input_mask', 'segment_ids'%5D%0A    inputs = %5B%5D%0A    %0A    examples = %5B%5D%0A
     example 
@@ -1744,16 +1744,19 @@
 (sources
+%5B0%5D
 )%0A    if
@@ -1779,16 +1779,20 @@
 :%0A      
+ 
+   
 raise Va
@@ -1834,160 +1834,430 @@
 .%22)%0A
-%0A
     
-# Predict tag IDs.%0A    keys = %5B'input_ids', 'input_mask', 'segment_ids'%5D%0A    out = self._predictor(%7Bkey: %5Bexample.features%5Bkey%5D%5D for key in keys%7D)%0A
+inputs = %7Bkey: %5Bexample.features%5Bkey%5D%5D for key in keys%7D%0A    examples.append(example)%0A    %0A    for source in sources%5B1:%5D:%0A        example = self._example_builder.build_bert_example(source)%0A        examples.append(example)%0A        for key in keys:%0A            inputs%5Bkey%5D.append(example.features%5Bkey%5D)%0A%0A    outputs = self._predictor(inputs)%5B'pred'%5D%0A    %0A    predictions = %5B%5D%0A    for i, output in enumerate(outputs):%0A    
     
@@ -2279,19 +2279,11 @@
  out
-%5B'pred'%5D%5B0%5D
+put
 .tol
@@ -2288,16 +2288,20 @@
 olist()%0A
+    
     # Re
@@ -2314,31 +2314,39 @@
 output.%0A    
+  
+  
 example
+s%5Bi%5D
 .features%5B'l
@@ -2377,48 +2377,8 @@
     
-# Mask out the begin and the end token.%0A
     
@@ -2380,24 +2380,28 @@
      example
+s%5Bi%5D
 .features%5B'l
@@ -2414,16 +2414,28 @@
 mask'%5D =
+ %5C%0A         
  %5B0%5D + %5B
@@ -2470,16 +2470,20 @@
 ) + %5B0%5D%0A
+    
     labe
@@ -2545,16 +2545,20 @@
  example
+s%5Bi%5D
 .get_tok
@@ -2577,28 +2577,52 @@
     
+  
-%5D%0A    return 
+  %5D%0A        predictions.append(
 example
+s%5Bi%5D
 .edi
@@ -2653,9 +2653,42 @@
 (labels)
+)%0A        %0A    return predictions
 %0A

Update predict_utils.py