Transformer model sentiment analysis example

examples/mnist/implementations/models/t2t.py

@@ -35,6 +35,6 @@ def transform_tensorflow(features, labels, model_config):
     features["inputs"] = tf.reshape(features["image_pixels"], hparams["input_shape"])
 
     # t2t expects this key and dimensionality
-    features["targets"] = tf.expand_dims(labels, 0)
+    features["targets"] = tf.expand_dims(tf.expand_dims(labels, -1), -1)
 
     return features, labels

examples/reviews/implementations/aggregators/vocab.py

@@ -2,7 +2,6 @@ def aggregate_spark(data, columns, args):
     import pyspark.sql.functions as F
     from pyspark.ml.feature import StopWordsRemover, RegexTokenizer
 
-    input_data = data.withColumn(columns["col"], F.lower(F.col(columns["col"])))
     regexTokenizer = RegexTokenizer(inputCol=columns["col"], outputCol="token_list", pattern="\\W")
     regexTokenized = regexTokenizer.transform(data)
 
@@ -19,6 +18,7 @@ def aggregate_spark(data, columns, args):
     )
 
     vocab = [row["word"] for row in vocab_rows]
-    reverse_dict = {word: idx + len(args["reserved_indices"]) for idx, word in enumerate(vocab)}
-
-    return {**reverse_dict, **args["reserved_indices"]}
+    reverse_dict = {word: 2 + idx for idx, word in enumerate(vocab)}
+    reverse_dict["<PAD>"] = 0
+    reverse_dict["<UNKNOWN>"] = 1
+    return {**reverse_dict}

examples/reviews/implementations/models/t2t_transformer.py

@@ -0,0 +1,59 @@
+import tensorflow as tf
+from tensor2tensor.utils import trainer_lib
+from tensor2tensor import models  # pylint: disable=unused-import
+from tensor2tensor import problems  # pylint: disable=unused-import
+from tensor2tensor.data_generators import problem_hparams
+from tensor2tensor.utils import registry
+from tensor2tensor.utils import metrics
+from tensor2tensor.data_generators import imdb
+from tensor2tensor.data_generators import text_encoder
+
+
+def create_estimator(run_config, model_config):
+    # t2t expects these keys in run_config
+    run_config.data_parallelism = None
+    run_config.t2t_device_info = {"num_async_replicas": 1}
+
+    hparams = trainer_lib.create_hparams("transformer_base_single_gpu")
+
+    problem = SentimentIMDBCortex(list(model_config["aggregates"]["reviews_vocab"]))
+    p_hparams = problem.get_hparams(hparams)
+    hparams.problem = problem
+    hparams.problem_hparams = p_hparams
+
+    # only want ACC
+    problem.eval_metrics = lambda: [metrics.Metrics.ACC]
+
+    # t2t expects this key
+    hparams.warm_start_from = None
+
+    estimator = trainer_lib.create_estimator("transformer", hparams, run_config)
+    return estimator
+
+
+def transform_tensorflow(features, labels, model_config):
+    max_length = model_config["aggregates"]["max_review_length"]
+
+    features["inputs"] = tf.expand_dims(
+        tf.expand_dims(tf.reshape(features["embedding_input"], [max_length]), -1), -1
+    )
+
+    features["targets"] = tf.expand_dims(tf.expand_dims(labels, -1), -1)
+
+    return features, labels
+
+
+class SentimentIMDBCortex(imdb.SentimentIMDB):
+    """IMDB sentiment classification, with an in-memory vocab"""
+
+    def __init__(self, vocab_list):
+        super().__init__()
+        self.vocab = vocab_list
+
+    def feature_encoders(self, data_dir):
+        encoder = text_encoder.TokenTextEncoder(vocab_filename=None, vocab_list=self.vocab)
+
+        return {
+            "inputs": encoder,
+            "targets": text_encoder.ClassLabelEncoder(self.class_labels(data_dir)),
+        }

examples/reviews/implementations/transformers/tokenize_string_to_int.py

@@ -1,26 +1,28 @@
 import re
+import nltk
+from nltk.corpus import stopwords
+
+nltk.download("stopwords")
 
 non_word = re.compile("\\W")
 
 
 def transform_python(sample, args):
     text = sample["col"].lower()
     token_index_list = []
-
-    reverse_vocab = args["vocab"]
-    stop_words = args["stop_words"]
-    reserved_indices = args["reserved_indices"]
+    vocab = args["vocab"]
+    stop_words = set(stopwords.words("english"))
 
     for token in non_word.split(text):
         if len(token) == 0:
             continue
         if token in stop_words:
             continue
-        token_index_list.append(reverse_vocab.get(token, reserved_indices["<UNKNOWN>"]))
+        token_index_list.append(vocab.get(token, 1))
         if len(token_index_list) == args["max_len"]:
             break
 
     for i in range(args["max_len"] - len(token_index_list)):
-        token_index_list.append(reserved_indices["<PAD>"])
+        token_index_list.append(0)
 
     return token_index_list

examples/reviews/requirements.txt

@@ -0,0 +1,2 @@
+tensor2tensor==1.10.0
+nltk==3.4

examples/reviews/resources/aggregates.yaml

@@ -6,7 +6,6 @@
       col: review
     args:
       vocab_size: 10000
-      reserved_indices: reserved_indices
 
 - kind: aggregate
   name: max_review_length

examples/reviews/resources/aggregators.yaml

@@ -6,7 +6,6 @@
       col: STRING_COLUMN
     args:
       vocab_size: INT
-      reserved_indices: {STRING: INT}
 
 - kind: aggregator
   name: max_length

examples/reviews/resources/apis.yaml

@@ -9,3 +9,9 @@
   model_name: sentiment_linear
   compute:
     replicas: 1
+
+- kind: api
+  name: sentiment-t2t
+  model_name: t2t_transformer
+  compute:
+    replicas: 1

examples/reviews/resources/models.yaml

@@ -30,3 +30,24 @@
   training:
     batch_size: 64
     num_steps: 5000
+
+- kind: model
+  name: t2t_transformer
+  type: classification
+  target_column: label_indexed
+  feature_columns:
+    - embedding_input
+  aggregates:
+    - max_review_length
+    - reviews_vocab
+  prediction_key: outputs
+  data_partition_ratio:
+    training: 0.8
+    evaluation: 0.2
+  training:
+    batch_size: 1
+    num_steps: 1
+  evaluation:
+    start_delay_secs: 1
+  compute:
+    gpu: 1

examples/reviews/resources/transformed_columns.yaml

@@ -6,9 +6,7 @@
       col: review
     args:
       max_len: max_review_length
-      stop_words: english_stop_words
       vocab: reviews_vocab
-      reserved_indices: reserved_indices
 
 - kind: transformed_column
   name: label_indexed

examples/reviews/resources/transformers.yaml

@@ -6,6 +6,4 @@
       col: STRING_COLUMN
     args:
       max_len: INT
-      stop_words: {STRING: BOOL}
       vocab: {STRING: INT}
-      reserved_indices: {STRING: INT}