txtai-simple-test.py

import time

from txtai.pipeline import Labels

# Create labels instances per model type
tflabels = Labels(train("microsoft/xtremedistil-l6-h384-uncased", ds["train"], logging_steps=2000), dynamic=False)

thlabels = Labels((model, tokenizer), dynamic=False)
thlabels.pipeline.model.config.id2label = {0: "sadness", 1: "joy", 2: "love", 3: "anger", 4: "fear", 5: "surprise"}

sklabels = Labels((skmodel, tokenize), dynamic=False)
sklabels.pipeline.model.config.id2label = {0: "sadness", 1: "joy", 2: "love", 3: "anger", 4: "fear", 5: "surprise"}

# Test inputs
inputs = ds["test"]["text"]
print("Testing speed of %d items" % len(inputs))

start = time.time()
tflabels(inputs)
print("Transformers time =", time.time() - start)

start = time.time()
thlabels(inputs)
print("PyTorch time =", time.time() - start)

start = time.time()
sklabels(inputs, multilabel=None)
print("TF-IDF + Logistic Regression time =", time.time() - start)