Skip to content

Instantly share code, notes, and snippets.

@barrucadu barrucadu/dcg.py
Created Sep 13, 2019

Embed
What would you like to do?
#!/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'))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.