|
---
|
|
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)
|
|
|
|
```python
|
|
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
|
|
|
|
```
|
|
|
