I'm generating pairwise sentence order probabilities in the following way:
import itertools
import random
import numpy as np
import tensorflow as tf
from transformers import BertTokenizer, TFBertForNextSentencePrediction
np.set_printoptions(suppress=True)
cache_dir = '/path/to/cache/dir'
pretrained_weights = 'bert-base-multilingual-cased'
tokenizer = BertTokenizer.from_pretrained(
pretrained_weights,
cache_dir=cache_dir,
)
model = TFBertForNextSentencePrediction.from_pretrained(
pretrained_weights,
cache_dir=cache_dir,
)
sentences = """
In “The Necklace” by Guy de Maupassant, the main character, Mathilde, has always dreamed of
being an aristocrat but lives in poverty.
Embarrassed about her lack of fine possessions, she borrows a necklace from a wealthy
friend but loses it.
The story is known for its subversive and influential twist ending
"""
sentences = [s.strip() for s in sentences.strip().split('.')]
random.shuffle(sentences)
print(sentences)
pairs = itertools.permutations(sentences, 2)
encoded = tokenizer.batch_encode_plus(pairs, return_tensors='np', padding=True)
outputs = model(encoded)
probs = tf.keras.activations.softmax(outputs[0])
for i, s in enumerate(sentences, 1):
print(f's{i}: {s}')
for s, prob in zip(itertools.permutations(['s1', 's2', 's3'], 2), probs):
print(s, prob)
I'm not sure how to interpret the resulting probabilities to generate the ordered sentences.
s1: Embarrassed about her lack of fine possessions, she borrows a necklace from a wealthy
friend but loses it
s2: In “The Necklace” by Guy de Maupassant, the main character, Mathilde, has always dreamed of
being an aristocrat but lives in poverty
s3: The story is known for its subversive and influential twist ending
('s1', 's2') tf.Tensor([0.9987061 0.00129389], shape=(2,), dtype=float32)
('s1', 's3') tf.Tensor([0.9514299 0.04857007], shape=(2,), dtype=float32)
('s2', 's1') tf.Tensor([0.9994491 0.00055089], shape=(2,), dtype=float32)
('s2', 's3') tf.Tensor([0.94130975 0.05869029], shape=(2,), dtype=float32)
('s3', 's1') tf.Tensor([0.15520796 0.84479207], shape=(2,), dtype=float32)
('s3', 's2') tf.Tensor([0.98460925 0.01539072], shape=(2,), dtype=float32)
Update: I created this hacky solution based on the one I found here (not sure about its accuracy) which orders sentences however, due to calculating the probabilities for the cartesian product of sentences, the total items to predict increases by a square factor ex: for a set of 88 sentences the total pairs would be 88 * 88 = 7744 which won't scale nicely. The inference speed wouldn't be much of a problem when done on gpu but still, better suggestions of achieving the same result are welcome.
class HashableDict(dict): # Passed to `tf.keras.Model.predict` to enable batching
def __hash__(self):
return hash(tuple(self.items()))
def create_correlation_matrix(sentences, tokenizer, model, **kwargs):
np.set_printoptions(suppress=True)
pairs = itertools.product(sentences, repeat=2)
encoded = tokenizer.batch_encode_plus(pairs, return_tensors='np', padding=True)
logits = model.predict(HashableDict(**encoded), **kwargs)
probs = tf.keras.activations.softmax(tf.convert_to_tensor(logits[0]))
size = len(sentences)
return probs[:, 0].numpy().reshape(size, size)
def reorder_sentences(sentences, tokenizer, model, **kwargs):
ordered = []
correlation_matrix = create_correlation_matrix(
sentences, tokenizer, model, **kwargs
)
idx = np.unravel_index(
np.argmax(correlation_matrix, axis=None), correlation_matrix.shape
)
while correlation_matrix.any():
x_idx = idx[1]
correlation_matrix[idx[0], :] = 0
correlation_matrix[:, idx[0]] = 0
ordered.append(idx[0])
idx = np.unravel_index(
np.argmax(correlation_matrix[x_idx, :], axis=None),
correlation_matrix[x_idx, :].shape,
)
idx = (x_idx, idx[0])
return ordered
if __name__ == '__main__':
cache_dir = '/path/to/cache/dir'
pretrained_weights = 'bert-base-multilingual-cased'
tok = BertTokenizer.from_pretrained(
pretrained_weights,
cache_dir=cache_dir,
)
m = TFBertForNextSentencePrediction.from_pretrained(
pretrained_weights,
cache_dir=cache_dir,
)
s = """
In “The Necklace” by Guy de Maupassant, the main character, Mathilde, has always dreamed of
being an aristocrat but lives in poverty.
Embarrassed about her lack of fine possessions, she borrows a necklace from a wealthy
friend but loses it.
The story is known for its subversive and influential twist ending
"""
s = [ss.strip() for ss in s.strip().split('.')]
print(s)
random.shuffle(s)
print(s)
ordering = reorder_sentences(s, tok, m, verbose=True, batch_size=8)
reordered_sentences = [s[idx] for idx in ordering]
print(ordering, reordered_sentences)
results in:
2/2 [==============================] - 9s 198ms/step
[1, 0, 2] ['In “The Necklace” by Guy de Maupassant, the main character, Mathilde, has always dreamed of \n being an aristocrat but lives in poverty', 'Embarrassed about her lack of fine possessions, she borrows a necklace from a wealthy \n friend but loses it', 'The story is known for its subversive and influential twist ending']