How to order sentences based on pairwise probabilities?

Viewed 74

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']
1 Answers

The output probabilities represent how likely it is that the first sentence is followed by the second one.

The two probabilities sum to 1 since this is a binary problem, and they represent the probability of answering "no" to the question (first value) and the probability of answering "yes" (second value).

For example in among your 3 sentences only s3 is predicted likely to be followed by s1, with probability 84%. All the other pairs are predicted as negative.

Following these results, you could obtain a partial graph to build the possible ordered sequences. There are probably several options to do that, for example you could start from the highest probabilities, or pick the most likely following sentence of each sentence.


Edit: design idea for obtaining the most likely order.

Note that this an open design question, I'm not aware of any standard method for this. This is probably a research question, so there might be research papers which address it. Below is how I would attempt to do it.

  • Represent the full graph of all possible transitions from one sentence to another one. Sentences are vertices and the directed arcs connect two sentences, with the probability as a weight. In this part it would probably be useful to discard all the arcs where the positive probability is less than some threshold $t$, say for example 0.1.
  • Then we need to find a (vertex-disjoint) path cover which maximizes the product of the weights/probabilities. It's a complex problem and it's probably not doable with an exact method (this is definitely above my paygrade!). A genetic algorithm would probably work. I could imagine an even more simplified solution as follows:
    • select a set of candidate first sentences among the ones which have no edge directed toward them, or only with low probability.
    • for each candidate first sentence, follow the path by always taking the most likely next sentence.
    • discard any path which doesn't connect all the sentences. In some cases there might be no solution, so there would have to be some workaround for this.
    • Finally pick the path which maximizes the product of the probabilities (tip: sum over the log of the probability)

The selected path gives the ordered sequence of sentences.

Related