test bigram counters with discontinuous bigrams and repeated words verify that two sequences of ngram association values are within epsilon of each other test bigram counters with discontinuous bigrams and repeated words verify that two sequences of n gram association values are within _epsilon of each other
from nltk.collocations import BigramCollocationFinder
from nltk.metrics import BigramAssocMeasures



_EPSILON = 1e-8
SENT = "this this is is a a test test".split()


def close_enough(x, y):
    

    return all(abs(x1[1] - y1[1]) <= _EPSILON for x1, y1 in zip(x, y))


def test_bigram2():
    b = BigramCollocationFinder.from_words(SENT)

    assert sorted(b.ngram_fd.items()) == [
        (("a", "a"), 1),
        (("a", "test"), 1),
        (("is", "a"), 1),
        (("is", "is"), 1),
        (("test", "test"), 1),
        (("this", "is"), 1),
        (("this", "this"), 1),
    ]
    assert sorted(b.word_fd.items()) == [("a", 2), ("is", 2), ("test", 2), ("this", 2)]

    assert len(SENT) == sum(b.word_fd.values()) == sum(b.ngram_fd.values()) + 1
    assert close_enough(
        sorted(b.score_ngrams(BigramAssocMeasures.pmi)),
        [
            (("a", "a"), 1.0),
            (("a", "test"), 1.0),
            (("is", "a"), 1.0),
            (("is", "is"), 1.0),
            (("test", "test"), 1.0),
            (("this", "is"), 1.0),
            (("this", "this"), 1.0),
        ],
    )


def test_bigram3():
    b = BigramCollocationFinder.from_words(SENT, window_size=3)
    assert sorted(b.ngram_fd.items()) == sorted(
        [
            (("a", "test"), 3),
            (("is", "a"), 3),
            (("this", "is"), 3),
            (("a", "a"), 1),
            (("is", "is"), 1),
            (("test", "test"), 1),
            (("this", "this"), 1),
        ]
    )

    assert sorted(b.word_fd.items()) == sorted(
        [("a", 2), ("is", 2), ("test", 2), ("this", 2)]
    )

    assert (
        len(SENT) == sum(b.word_fd.values()) == (sum(b.ngram_fd.values()) + 2 + 1) / 2.0
    )
    assert close_enough(
        sorted(b.score_ngrams(BigramAssocMeasures.pmi)),
        sorted(
            [
                (("a", "test"), 1.584962500721156),
                (("is", "a"), 1.584962500721156),
                (("this", "is"), 1.584962500721156),
                (("a", "a"), 0.0),
                (("is", "is"), 0.0),
                (("test", "test"), 0.0),
                (("this", "this"), 0.0),
            ]
        ),
    )


def test_bigram5():
    b = BigramCollocationFinder.from_words(SENT, window_size=5)
    assert sorted(b.ngram_fd.items()) == sorted(
        [
            (("a", "test"), 4),
            (("is", "a"), 4),
            (("this", "is"), 4),
            (("is", "test"), 3),
            (("this", "a"), 3),
            (("a", "a"), 1),
            (("is", "is"), 1),
            (("test", "test"), 1),
            (("this", "this"), 1),
        ]
    )
    assert sorted(b.word_fd.items()) == sorted(
        [("a", 2), ("is", 2), ("test", 2), ("this", 2)]
    )
    n_word_fd = sum(b.word_fd.values())
    n_ngram_fd = (sum(b.ngram_fd.values()) + 4 + 3 + 2 + 1) / 4.0
    assert len(SENT) == n_word_fd == n_ngram_fd
    assert close_enough(
        sorted(b.score_ngrams(BigramAssocMeasures.pmi)),
        sorted(
            [
                (("a", "test"), 1.0),
                (("is", "a"), 1.0),
                (("this", "is"), 1.0),
                (("is", "test"), 0.5849625007211562),
                (("this", "a"), 0.5849625007211562),
                (("a", "a"), -1.0),
                (("is", "is"), -1.0),
                (("test", "test"), -1.0),
                (("this", "this"), -1.0),
            ]
        ),
    )