dcg.py

#!/usr/bin/env nix-shell
#! nix-shell -i python3 -p "python3.withPackages (ps: [ps.requests ps.tabulate])"

import csv
import math
import os
import requests

from tabulate import tabulate

SEARCH_API = os.getenv('SEARCH_API', 'https://www.gov.uk/api/search.json')
AB_TEST_NAME = os.getenv('AB_TEST_NAME', 'search_cluster_query')
AB_TEST_VARIANTS = os.getenv('AB_TEST_VARIANTS', 'A,B').split(',')
DATA = os.getenv('DATA', 'relevancy.txt')
SEPARATE_TIES = os.getenv('SEPARATE_TIES', '') != ''

queries = {}

with open(DATA) as csvfile:
    rows = csv.reader(csvfile)
    next(rows)
    for row in rows:
        query = row[0].strip()
        rating = row[1].strip().lower()
        link = row[2].strip()
        if rating == 'relevant':
            rating = 3
        elif rating == 'near':
            rating = 2
        elif rating == 'misplaced':
            rating = 1
        elif rating in ['irrelevant', '0']:
            continue
        else:
            rating = int(row[1])
        judgements = queries.get(query, {})
        judgements[link] = rating
        queries[query] = judgements

results = []
winner = {}
for query, judgements in queries.items():
    print(query)
    row = [query]
    scores = []
    best = None
    best_score = None
    worst_score = None
    for variant in AB_TEST_VARIANTS:
        r = requests.get(SEARCH_API, params={'q': query, 'ab_tests': f'{AB_TEST_NAME}:{variant}'})
        # https://en.wikipedia.org/wiki/Discounted_cumulative_gain
        score = 0
        try:
            for rank, result in zip(range(1, 21), r.json()['results']):
                judgement = judgements.get(result['link'], 0)
                score += (2**judgement - 1) / math.log (rank + 1)
        except:
            print(query, variant)
            print(r.json())
        if worst_score is None or score < worst_score:
            worst_score = score
        if best is None or best_score < score:
            best = [variant]
            best_score = score
        elif best_score == score:
            best.append(variant)
        scores.append(score)
    row.append(','.join(best))
    row.append('!' if best_score * 0.95 > worst_score else '')

    if len(best) > 1 and SEPARATE_TIES:
        bests = ','.join(best)
        winner[bests] = winner.get(bests, 0) + 1
    else:
        for b in best:
            winner[b] = winner.get(b, 0) + 1
    row.extend(scores)
    results.append(row)

headers = ['query', 'best', '!']
headers.extend(AB_TEST_VARIANTS)
print(tabulate(results, headers=headers, tablefmt='presto'))
print()
print(tabulate(sorted(winner.items(), key=lambda row: row[1]), headers=['variant', 'wins'], tablefmt='presto'))