Update

README.md

@@ -17,68 +17,131 @@ This repository contains a small ELECTRA discriminator finetuned on a corpus of
 
 ```python
 import math
+import numpy as np
 
 import tensorflow as tf
-from datasets import Dataset, ClassLabel, Features, Value
-from transformers import TFAutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, create_optimizer
+from datasets import load_metric, Dataset, DatasetDict
+from transformers import TFAutoModel, TFAutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, create_optimizer
+from transformers.keras_callbacks import KerasMetricCallback
 
 # This example shows how this model can be used:
-#  you should finetune the model of your specific corpus if commands, bogger than this
+#  you should finetune the model of your specific corpus if commands, bigger than this
 dict_train = {
-    "idx": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15"],
-    "sentence": ["e", "get pen", "drop book", "x paper", "i", "south", "get paper", "drop pen", "x book", "inventory",
-                 "n", "get book", "drop paper", "examine Pen", "inv", "w"],
-    "label": ["v01835496", "v01214265", "v01977701", "v02131279", "v02472495", "v01835496", "v01214265", "v01977701",
-              "v02131279", "v02472495", "v01835496", "v01214265", "v01977701", "v02131279", "v02472495", "v01835496"]
+    "idx": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17", "18",
+            "19", "20"],
+    "sentence": ["e", "get pen", "drop book", "x paper", "i", "south", "get paper", "drop the pen", "x book",
+                 "inventory", "n", "get the book", "drop paper", "look at Pen", "inv", "g", "s", "get sandwich",
+                 "drop sandwich", "x sandwich", "agin"],
+    "label": ["travel.v.01", "take.v.04", "drop.v.01", "examine.v.02", "inventory.v.01", "travel.v.01", "take.v.04",
+              "drop.v.01", "examine.v.02", "inventory.v.01", "travel.v.01", "take.v.04", "drop.v.01", "examine.v.02",
+              "inventory.v.01", "repeat.v.01", "travel.v.01", "take.v.04", "drop.v.01", "examine.v.02", "repeat.v.01"]
+}
+dict_val = {
+    "idx": ["0", "1", "2", "3", "4", "5"],
+    "sentence": ["w", "get shield", "drop sword", "x spikes", "i", "repeat"],
+    "label": ["travel.v.01", "take.v.04", "drop.v.01", "examine.v.02", "inventory.v.01", "repeat.v.01"]
 }
 
-num_labels = len(set(dict_train["label"]))
-features = Features({'idx': Value('uint32'), 'sentence': Value('string'),
-                     'label': ClassLabel(names=list(set(dict_train["label"])))})
-
-raw_train_dataset = Dataset.from_dict(dict_train, features=features)
-
-discriminator = TFAutoModelForSequenceClassification.from_pretrained("Aureliano/electra-if", num_labels=num_labels)
+raw_train_dataset = Dataset.from_dict(dict_train)
+raw_val_dataset = Dataset.from_dict(dict_val)
+raw_dataset = DatasetDict()
+raw_dataset["train"] = raw_train_dataset
+raw_dataset["val"] = raw_val_dataset
+raw_dataset = raw_dataset.class_encode_column("label")
+print(raw_dataset)
+print(raw_dataset["train"].features)
+print(raw_dataset["val"].features)
+print(raw_dataset["train"][1])
+label2id = {}
+id2label = {}
+for i, l in enumerate(raw_dataset["train"].features["label"].names):
+    label2id[l] = i
+    id2label[i] = l
+
+discriminator = TFAutoModelForSequenceClassification.from_pretrained("Aureliano/electra-if",
+                                                                     label2id=label2id,
+                                                                     id2label=id2label)
 tokenizer = AutoTokenizer.from_pretrained("Aureliano/electra-if")
 
 tokenize_function = lambda example: tokenizer(example["sentence"], truncation=True)
 
-pre_tokenizer_columns = set(raw_train_dataset.features)
-train_dataset = raw_train_dataset.map(tokenize_function, batched=True)
-tokenizer_columns = list(set(train_dataset.features) - pre_tokenizer_columns)
+pre_tokenizer_columns = set(raw_dataset["train"].features)
+encoded_dataset = raw_dataset.map(tokenize_function, batched=True)
+tokenizer_columns = list(set(encoded_dataset["train"].features) - pre_tokenizer_columns)
 
 data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")
 
-batch_size = 16
-tf_train_dataset = train_dataset.to_tf_dataset(
+batch_size = len(encoded_dataset["train"])
+tf_train_dataset = encoded_dataset["train"].to_tf_dataset(
     columns=tokenizer_columns,
     label_cols=["labels"],
     shuffle=True,
     batch_size=batch_size,
     collate_fn=data_collator
 )
+tf_validation_dataset = encoded_dataset["val"].to_tf_dataset(
+    columns=tokenizer_columns,
+    label_cols=["labels"],
+    shuffle=False,
+    batch_size=batch_size,
+    collate_fn=data_collator
+)
 
 loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
-num_epochs = 100
-batches_per_epoch = math.ceil(len(train_dataset) / batch_size)
+num_epochs = 20
+batches_per_epoch = math.ceil(len(encoded_dataset["train"]) / batch_size)
 total_train_steps = int(batches_per_epoch * num_epochs)
 
 optimizer, schedule = create_optimizer(
-    init_lr=1e-5, num_warmup_steps=1, num_train_steps=total_train_steps
+    init_lr=5e-5, num_warmup_steps=total_train_steps // 5, num_train_steps=total_train_steps
 )
 
-discriminator.compile(optimizer=optimizer, loss=loss)
+metric = load_metric("accuracy")
+
+
+def compute_metrics(eval_predictions):
+    logits, labels = eval_predictions
+    predictions = np.argmax(logits, axis=-1)
+    return metric.compute(predictions=predictions, references=labels)
+
+
+metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=tf_validation_dataset)
+callbacks = [metric_callback]
+
+discriminator.compile(optimizer=optimizer, loss=loss, metrics=["sparse_categorical_accuracy"])
 discriminator.fit(
     tf_train_dataset,
-    epochs=num_epochs
+    epochs=num_epochs,
+    validation_data=tf_validation_dataset,
+    callbacks=callbacks
 )
 
+print("Evaluate on test data")
+results = discriminator.evaluate(tf_validation_dataset)
+print("test loss, test acc:", results)
+
+text = "i"
+encoded_input = tokenizer(text, return_tensors='tf')
+output = discriminator(encoded_input)
+prediction = tf.nn.softmax(output["logits"][0], -1)
+label = id2label[tf.math.argmax(prediction).numpy()]
+print("\n", text, ":", label,
+      "\n")  # ideally 'inventory.v.01' (-> "make or include in an itemized record or report"), but probably only with a better finetuning dataset
+
 text = "get lamp"
 encoded_input = tokenizer(text, return_tensors='tf')
 output = discriminator(encoded_input)
 prediction = tf.nn.softmax(output["logits"][0], -1)
-label = dict_train["label"][tf.math.argmax(prediction)]
-print(text, ":", label)
-# ideally [v01214265 -> take.v.04 -> "get into one's hands, take physically"], but probably only with a better dataset
+label = id2label[tf.math.argmax(prediction).numpy()]
+print("\n", text, ":", label,
+      "\n")  # ideally 'take.v.04' (-> "get into one's hands, take physically"), but probably only with a better finetuning dataset
+
+text = "w"
+encoded_input = tokenizer(text, return_tensors='tf')
+output = discriminator(encoded_input)
+prediction = tf.nn.softmax(output["logits"][0], -1)
+label = id2label[tf.math.argmax(prediction).numpy()]
+print("\n", text, ":", label,
+      "\n")  # ideally 'travel.v.01' (-> "change location; move, travel, or proceed, also metaphorically"), but probably only with a better finetuning dataset
 
 ```

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a613fb2aa3569bac1a4a5cb2be88706cdccd33a72dddfc78051353e6ec07cb46
+oid sha256:b585bc34677a9d35a1ef60bf1a531067fad58e9f4259dda25d63b4ae54632f93
 size 54011377

tf_model.h5

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:56a0bfbd98a2fb0c1a57ad0c8a4dd6234cc17b2e27a21b5c33971a5045d5ecd7
+oid sha256:4bf1f25a9684bfe9bb90e65a365b9b10bdd3b413841c98dba94b900139af7c90
 size 54198792