metadata
license: apache-2.0
How to use the discriminator in transformers on a custom dataset
(Heavily based on: https://github.com/huggingface/notebooks/blob/master/examples/text_classification-tf.ipynb)
import math
import tensorflow as tf
from datasets import Dataset, ClassLabel, Features, Value
from transformers import TFAutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, create_optimizer
# This example shows how this model can be used:
# you should finetune the model of your specific corpus if commands, bogger 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"]
}
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/distilbert-base-uncased-if", num_labels=num_labels)
tokenizer = AutoTokenizer.from_pretrained("Aureliano/distilbert-base-uncased-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)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")
batch_size = 16
tf_train_dataset = train_dataset.to_tf_dataset(
columns=tokenizer_columns,
label_cols=["labels"],
shuffle=True,
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)
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
)
discriminator.compile(optimizer=optimizer, loss=loss)
discriminator.fit(
tf_train_dataset,
epochs=num_epochs
)
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