<commit_before>from math import log


def _calc_dcg(judgments, indices):
    return sum([(2**j - 1) / log(i + 2, 2) for (i, j) in zip(indices, judgments)])


def compute_ndcg(iterable, acceptable_responses=frozenset([0, 1, 2, 3]), indices=None):
    """
    :param iterable: the iterable full of judgements
    :param acceptable_responses: a set containing all legal values
    :param indices: the indices of the list, in case we skip a ranking
    """
    if not indices:
        indices = range(len(iterable))
    dcg = _calc_dcg(iterable, indices)
    ideal_order = sorted(iterable, reverse=True)
    idcg = _calc_dcg(ideal_order, indices)
    try:
        ndcg = dcg/idcg
    except Exception:
        print "Uh oh! Error calculating"
        print iterable
        print dcg, "/", idcg
        ndcg = None
    return ndcg


def close_enough(a, b):
    if abs(a-b) < .0000001:
        return True
    return False


def test():
    lst = [3, 2, 3, 1, 2, 0]
    exp = 0.958112357102
    assert close_enough(compute_ndcg(lst), exp), "Uh oh! Expecting {} for {}".format(exp, lst)
    lst = [3, 2, 3, 1, 2, 1]
    exp = 0.959110289196
    assert close_enough(compute_ndcg(lst), exp), "Uh oh! Expecting {} for {}".format(exp, lst)
    lst2 = [i-1 for i in lst]
    exp = 0.947508362289
    assert close_enough(compute_ndcg(lst2, set([-1, 0, 2, 1])), exp), "Uh oh! Expecting {} for {}".format(exp, lst)


if __name__ == "__main__":
    test()<commit_msg>caveat about the correctness of the ndcg tests<commit_after>from math import log


def _calc_dcg(judgments, indices):
    return sum([(2**j - 1) / log(i + 2, 2) for (i, j) in zip(indices, judgments)])


def compute_ndcg(iterable, acceptable_responses=frozenset([0, 1, 2, 3]), indices=None):
    """
    :param iterable: the iterable full of judgements
    :param acceptable_responses: a set containing all legal values
    :param indices: the indices of the list, in case we skip a ranking
    """
    if not indices:
        indices = range(len(iterable))
    dcg = _calc_dcg(iterable, indices)
    ideal_order = sorted(iterable, reverse=True)
    idcg = _calc_dcg(ideal_order, indices)
    try:
        ndcg = dcg/idcg
    except Exception:
        print "Uh oh! Error calculating"
        print iterable
        print dcg, "/", idcg
        ndcg = None
    return ndcg


def close_enough(a, b):
    if abs(a-b) < .0000001:
        return True
    return False


def test():
    # TODO: these are asserting that the values are similar to what I got on Aug 10 2015
    # not any sort of ground truth or anything.
    lst = [3, 2, 3, 1, 2, 0]
    exp = 0.958112357102
    assert close_enough(compute_ndcg(lst), exp), "Uh oh! Expecting {} for {}".format(exp, lst)
    lst = [3, 2, 3, 1, 2, 1]
    exp = 0.959110289196
    assert close_enough(compute_ndcg(lst), exp), "Uh oh! Expecting {} for {}".format(exp, lst)
    lst2 = [i-1 for i in lst]
    exp = 0.947508362289
    assert close_enough(compute_ndcg(lst2, set([-1, 0, 2, 1])), exp), "Uh oh! Expecting {} for {}".format(exp, lst)


if __name__ == "__main__":
    test()
{'fraction_non_alphanumeric': 0.1262199089134678, 'fraction_numerical': 0.0494469746258946, 'mean_word_length': 3.00390625, 'pattern_counts': {'":': 2, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5245551', 'n_tokens_mistral': 1216, 'n_tokens_neox': 1076, 'n_words': 351}