class Tfidf:
def __init__(self):
self.config = Config("tfidf")
self.dataloader = Dataloader(self.config)
self.vectorizer = TfidfVectorizer()
def evaluate(self, metrics=None):
if metrics is None:
metrics = ["bleu", "rouge", "distinct1", "distinct2"]
# Fit to data
contexts, responses = self.dataloader.get_retrieval_candidates()
X = self.vectorizer.fit_transform(contexts)
# Evaluate on test data
t_contexts, t_responses = self.dataloader.get_test_data()
test_contexts = self.vectorizer.transform(t_contexts)
predictions = []
for context in test_contexts:
cosine_similarities = linear_kernel(context, X).flatten()
argsort = cosine_similarities.argsort()
predictions.append(responses[argsort[-1]])
pickle.dump((t_contexts, t_responses, predictions), open("Save/tfidf_test", "wb"))
results = {}
if "bleu" in metrics:
results["bleu"] = get_bleu(t_responses, predictions)
if "rouge" in metrics:
results["rouge-l"] = get_rouge(t_responses, predictions)
if "distinct1" in metrics:
results["distinct-1"] = get_distinct_1(predictions)
if "distinct2" in metrics:
results["distinct-2"] = get_distinct_2(predictions)
return results
class Trainer:
def __init__(self, model_name, config="auto"):
"""
Trainer class is used for managing the end-to-end training procedure for the model, through to evaluation.
:param model_name: string name of model
:param config: config object containing details of model and training hyper-parameters. Non-default parameters
can be added by including them as dictionary, .e.g config={"model_size": 256}
:raises ModelNotFoundError if the model_name given does not match the list of available models
To run the training process:
trainer = Trainer("baseline")
trainer.train()
# Some training happens
# ...
results = trainer.evaluate(["bleu", "rouge"])
# {"bleu": 0.67, "rouge-l": 0.5}
"""
if model_name not in ModelNotFoundError.model_names:
raise ModelNotFoundError()
if not os.path.isdir("Save"):
os.mkdir("Save")
self.config = Config(model_name)
self.tokenizer = self.config.tokenizer
if config != "auto":
assert type(config) is dict
for key, value in config.items():
self.config.__setattr__(key, value)
self.dataloader = Dataloader(self.config,
multitask=self.config.multitask)
self.transformer = Transformer(self.config)
opt = tf.keras.optimizers.Adam(learning_rate=self.config.learning_rate)
self.transformer.compile(optimizer=opt)
def train(self):
"""
Trains the model according to the details provided in config.
:return: None
"""
dataset = tf.data.Dataset.from_generator(
self.dataloader.generator, (np.int32, np.int32, np.int32))
trailing_losses = {
"generator": None,
"retrieval": None,
"reranker": None
}
validation_data = self.dataloader.get_validation_data()
validation_losses = []
for batch, inputs in enumerate(dataset.take(-1)):
losses = self.transformer.train_step(inputs)
if batch == 0:
for key, value in losses.items():
trailing_losses[key] = value
else:
for key, value in losses.items():
trailing_losses[
key] = self.config.momentum * trailing_losses[key] + (
1 - self.config.momentum) * value
# Reduce learning rate every DECAY_FREQ steps
if (batch + 1) % self.config.decay_freq == 0:
self.config.learning_rate *= self.config.decay_rate
self.transformer.optimizer.__setattr__(
"learning_rate", self.config.learning_rate)
if (batch + 1) % self.config.display_loss_freq == 0:
print("Step: ", batch + 1)
for key, value in trailing_losses.items():
if value is not None:
print(f"{key} loss: {np.round(value, 3)}")
print(self.transformer.predict())
self.transformer.encoder.save_weights(
f"Save/{self.config.model_name}_encoder.h5")
self.transformer.decoder.save_weights(
f"Save/{self.config.model_name}_decoder.h5")
if batch % self.config.validation_freq == 0:
losses = self.transformer.validation_loss(validation_data)
print("Step: ", batch + 1)
for key, value in losses.items():
print(f"{key} loss: {np.round(value, 3)}")
validation_losses.append(list(losses.values()))
if (batch + 1) == self.config.num_steps:
pickle.dump(
validation_losses,
open(f"Save/{self.config.model_name}_losses", "wb"))
break
def store_retrieval_candidates(self):
"""
Stores a pickle file of the hidden states of the retrieval candidates and their text representations for use in
inference.
:return: None
"""
contexts, responses = self.dataloader.get_retrieval_candidates()
chunk_size = 100
encoded_responses = []
for i in range(0, self.config.retrieval_candidates, chunk_size):
print(i)
try:
chunk_responses = responses[i:i + chunk_size]
# Encode candidates
candidates_encoded = self.tokenizer.encode_ids_with_bos_eos(
chunk_responses)
candidates = np.zeros([chunk_size, self.config.max_length])
for j, c in enumerate(candidates_encoded):
for k in range(
len(c) - 1
): # Leave off EOS, as Encoder was only trained on responses with BOS token
candidates[j, k] = c[k]
# Truncate
if k == (self.config.max_length - 1):
break
candidates = self.transformer.encoder.encode_responses(
tf.constant(candidates))
encoded_responses += (list(candidates.numpy()))
except IndexError as e:
print(e)
print("error")
encoded_responses = np.asarray(encoded_responses)
pickle.dump(encoded_responses, open("response_vectors", "wb"))
pickle.dump(responses, open("response_texts", "wb"))
def evaluate(self, metrics=None):
"""
Evaluates the model performance by first storing retrieval candidates if the model is a hybrid, then making
predictions against the test set, and finally measuring the score according to the automated metrics specified.
:param metrics: list of automated metrics. Options are "bleu", "rouge", "distinct1", "distinct2". All metrics
will be used if none are specified
:return: Dictionary of automated metrics and their associated scores
"""
# Check if retrieval candidates exist
if self.config.multitask:
if not (os.path.isfile("Save/response_vectors")
and os.path.isfile("Save/response_texts")):
self.store_retrieval_candidates()
# Check if model predictions exist
if not os.path.isfile(f"Save/{self.config.model_name}_test"):
contexts, responses = self.dataloader.get_test_data()
self.transformer.test(contexts, responses)
if metrics is None:
metrics = ["bleu", "rouge", "distinct1", "distinct2"]
_, responses, predictions = pickle.load(
open(f"Save/{self.config.model_name}_test", "rb"))
# Truncate responses to MAX_LENGTH tokens as this is what model made predictions against
responses = self.tokenizer.encode_ids(responses)
new_responses = []
for r in responses:
new_responses.append(r[:self.config.max_length])
responses = self.tokenizer.decode_ids(new_responses)
results = {}
if "bleu" in metrics:
results["bleu"] = get_bleu(responses, predictions)
if "rouge" in metrics:
results["rouge-l"] = get_rouge(responses, predictions)
if "distinct1" in metrics:
results["distinct-1"] = get_distinct_1(predictions)
if "distinct2" in metrics:
results["distinct-2"] = get_distinct_2(predictions)
return results