| import json |
| import os |
| from contextlib import contextmanager |
|
|
| import psycopg2 |
| from psycopg2.extras import RealDictCursor |
|
|
| DATABASE_URL = os.environ['DATABASE_URL'] |
|
|
| TABLE_ANNOTATORS = 'annotators_pp2' |
| TABLE_COMBINATIONS = 'combinations_pp2' |
| TABLE_ASSIGNMENTS = 'combo_assignments_pp2' |
| TABLE_ANNOTATIONS = 'annotations_pp2' |
|
|
|
|
| @contextmanager |
| def get_conn(): |
| conn = psycopg2.connect(DATABASE_URL) |
| try: |
| yield conn |
| conn.commit() |
| except Exception: |
| conn.rollback() |
| raise |
| finally: |
| conn.close() |
|
|
|
|
| def create_tables() -> None: |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| cur.execute( |
| f''' |
| CREATE EXTENSION IF NOT EXISTS pgcrypto; |
| |
| CREATE TABLE IF NOT EXISTS {TABLE_ANNOTATORS} ( |
| id UUID PRIMARY KEY DEFAULT gen_random_uuid(), |
| code INT UNIQUE NOT NULL CHECK (code BETWEEN 1 AND 8), |
| created_at TIMESTAMPTZ DEFAULT now() |
| ); |
| |
| CREATE TABLE IF NOT EXISTS {TABLE_COMBINATIONS} ( |
| id TEXT PRIMARY KEY, |
| base_id TEXT NOT NULL, |
| seed_id TEXT, |
| level INT, |
| class_tag TEXT, |
| ratio_type TEXT, |
| matchup_key TEXT NOT NULL, |
| model_left TEXT NOT NULL, |
| model_right TEXT NOT NULL, |
| n_total INT, |
| n_humans INT, |
| n_objects INT, |
| prompt TEXT, |
| prompt_zh TEXT, |
| humans JSONB, |
| objects JSONB |
| ); |
| |
| CREATE TABLE IF NOT EXISTS {TABLE_ASSIGNMENTS} ( |
| combo_id TEXT PRIMARY KEY REFERENCES {TABLE_COMBINATIONS}(id) ON DELETE CASCADE, |
| group_key TEXT, |
| annotator_a_code INT NOT NULL CHECK (annotator_a_code BETWEEN 1 AND 8), |
| annotator_b_code INT NOT NULL CHECK (annotator_b_code BETWEEN 1 AND 8), |
| created_at TIMESTAMPTZ DEFAULT now() |
| ); |
| |
| CREATE TABLE IF NOT EXISTS {TABLE_ANNOTATIONS} ( |
| id UUID PRIMARY KEY DEFAULT gen_random_uuid(), |
| annotator_id UUID NOT NULL REFERENCES {TABLE_ANNOTATORS}(id), |
| combo_id TEXT NOT NULL REFERENCES {TABLE_COMBINATIONS}(id), |
| model_a TEXT NOT NULL, |
| model_b TEXT NOT NULL, |
| a_existence INT CHECK (a_existence IN (0, 1)), |
| a_appearance INT CHECK (a_appearance IN (0, 1)), |
| a_interaction INT CHECK (a_interaction IN (0, 1)), |
| b_existence INT CHECK (b_existence IN (0, 1)), |
| b_appearance INT CHECK (b_appearance IN (0, 1)), |
| b_interaction INT CHECK (b_interaction IN (0, 1)), |
| preference TEXT CHECK (preference IN ('A', 'B')), |
| prompt_ilogical BOOLEAN NOT NULL DEFAULT FALSE, |
| created_at TIMESTAMPTZ DEFAULT now(), |
| UNIQUE (annotator_id, combo_id) |
| ); |
| ''' |
| ) |
| for code in range(1, 9): |
| cur.execute( |
| f"INSERT INTO {TABLE_ANNOTATORS} (code) VALUES (%s) ON CONFLICT (code) DO NOTHING", |
| (code,), |
| ) |
|
|
|
|
| def clear_round_tables() -> None: |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| cur.execute(f'TRUNCATE TABLE {TABLE_ANNOTATIONS} CASCADE') |
| cur.execute(f'TRUNCATE TABLE {TABLE_ASSIGNMENTS} CASCADE') |
| cur.execute(f'TRUNCATE TABLE {TABLE_COMBINATIONS} CASCADE') |
|
|
|
|
| def get_annotator_id_by_code(code: int) -> str: |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| cur.execute(f'SELECT id FROM {TABLE_ANNOTATORS} WHERE code = %s', (code,)) |
| row = cur.fetchone() |
| if not row: |
| raise ValueError(f'annotator code not found: {code}') |
| return str(row[0]) |
|
|
|
|
| def get_candidate_combos(annotator_code: int, limit: int = 500) -> list[dict]: |
| with get_conn() as conn: |
| with conn.cursor(cursor_factory=RealDictCursor) as cur: |
| cur.execute( |
| f''' |
| SELECT c.* |
| FROM {TABLE_COMBINATIONS} c |
| JOIN {TABLE_ASSIGNMENTS} a ON a.combo_id = c.id |
| JOIN {TABLE_ANNOTATORS} an ON an.code = %s |
| WHERE (a.annotator_a_code = %s OR a.annotator_b_code = %s) |
| AND NOT EXISTS ( |
| SELECT 1 FROM {TABLE_ANNOTATIONS} ann |
| WHERE ann.combo_id = c.id AND ann.annotator_id = an.id |
| ) |
| ORDER BY c.seed_id, c.matchup_key, c.level DESC, c.id |
| LIMIT %s |
| ''', |
| (annotator_code, annotator_code, annotator_code, limit), |
| ) |
| return [dict(r) for r in cur.fetchall()] |
|
|
|
|
| def count_remaining(annotator_code: int) -> int: |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| cur.execute( |
| f''' |
| SELECT COUNT(*) |
| FROM {TABLE_COMBINATIONS} c |
| JOIN {TABLE_ASSIGNMENTS} a ON a.combo_id = c.id |
| JOIN {TABLE_ANNOTATORS} an ON an.code = %s |
| WHERE (a.annotator_a_code = %s OR a.annotator_b_code = %s) |
| AND NOT EXISTS ( |
| SELECT 1 FROM {TABLE_ANNOTATIONS} ann |
| WHERE ann.combo_id = c.id AND ann.annotator_id = an.id |
| ) |
| ''', |
| (annotator_code, annotator_code, annotator_code), |
| ) |
| return int(cur.fetchone()[0]) |
|
|
|
|
| def save_annotation_by_code(data: dict) -> None: |
| annotator_id = get_annotator_id_by_code(int(data['annotator_code'])) |
| payload = { |
| 'annotator_id': annotator_id, |
| 'combo_id': data['combo_id'], |
| 'model_a': data['model_a'], |
| 'model_b': data['model_b'], |
| 'a_existence': data['a_existence'], |
| 'a_appearance': data['a_appearance'], |
| 'a_interaction': data['a_interaction'], |
| 'b_existence': data['b_existence'], |
| 'b_appearance': data['b_appearance'], |
| 'b_interaction': data['b_interaction'], |
| 'preference': data['preference'], |
| 'prompt_ilogical': bool(data.get('prompt_ilogical', False)), |
| } |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| cur.execute( |
| f''' |
| INSERT INTO {TABLE_ANNOTATIONS} |
| (annotator_id, combo_id, model_a, model_b, |
| a_existence, a_appearance, a_interaction, |
| b_existence, b_appearance, b_interaction, |
| preference, prompt_ilogical) |
| VALUES |
| (%(annotator_id)s, %(combo_id)s, %(model_a)s, %(model_b)s, |
| %(a_existence)s, %(a_appearance)s, %(a_interaction)s, |
| %(b_existence)s, %(b_appearance)s, %(b_interaction)s, |
| %(preference)s, %(prompt_ilogical)s) |
| ON CONFLICT (annotator_id, combo_id) DO UPDATE SET |
| model_a = EXCLUDED.model_a, |
| model_b = EXCLUDED.model_b, |
| a_existence = EXCLUDED.a_existence, |
| a_appearance = EXCLUDED.a_appearance, |
| a_interaction = EXCLUDED.a_interaction, |
| b_existence = EXCLUDED.b_existence, |
| b_appearance = EXCLUDED.b_appearance, |
| b_interaction = EXCLUDED.b_interaction, |
| preference = EXCLUDED.preference, |
| prompt_ilogical = EXCLUDED.prompt_ilogical |
| ''', |
| payload, |
| ) |
|
|
|
|
| def get_progress(annotator_code: int) -> dict: |
| annotator_id = get_annotator_id_by_code(annotator_code) |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| cur.execute(f'SELECT COUNT(*) FROM {TABLE_ANNOTATIONS} WHERE annotator_id = %s', (annotator_id,)) |
| completed = int(cur.fetchone()[0]) |
| cur.execute( |
| f''' |
| SELECT COUNT(*) |
| FROM {TABLE_COMBINATIONS} c |
| JOIN {TABLE_ASSIGNMENTS} a ON a.combo_id = c.id |
| WHERE (a.annotator_a_code = %s OR a.annotator_b_code = %s) |
| ''', |
| (annotator_code, annotator_code), |
| ) |
| quota = int(cur.fetchone()[0]) |
| cur.execute( |
| f''' |
| SELECT COUNT(*) FROM ( |
| SELECT combo_id |
| FROM {TABLE_ANNOTATIONS} |
| GROUP BY combo_id |
| HAVING COUNT(DISTINCT annotator_id) >= 2 |
| ) s |
| ''' |
| ) |
| overall_completed = int(cur.fetchone()[0]) |
| cur.execute(f'SELECT COUNT(*) FROM {TABLE_COMBINATIONS}') |
| overall_total = int(cur.fetchone()[0]) |
| return { |
| 'completed': completed, |
| 'quota': quota, |
| 'overall_completed': overall_completed, |
| 'overall_total': overall_total, |
| } |
|
|
|
|
| def get_progress_board() -> dict: |
| board = [] |
| overall_completed = 0 |
| overall_total = 0 |
| for code in range(1, 9): |
| p = get_progress(code) |
| board.append({'code': code, 'completed': p['completed'], 'quota': p['quota']}) |
| overall_completed = p['overall_completed'] |
| overall_total = p['overall_total'] |
| return {'annotators': board, 'overall_completed': overall_completed, 'overall_total': overall_total} |
|
|
|
|
| def upsert_combinations(combos: list[dict]) -> None: |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| for c in combos: |
| cur.execute( |
| f''' |
| INSERT INTO {TABLE_COMBINATIONS} |
| (id, base_id, seed_id, level, class_tag, ratio_type, matchup_key, model_left, model_right, |
| n_total, n_humans, n_objects, prompt, prompt_zh, humans, objects) |
| VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) |
| ON CONFLICT (id) DO UPDATE SET |
| base_id = EXCLUDED.base_id, |
| seed_id = EXCLUDED.seed_id, |
| level = EXCLUDED.level, |
| class_tag = EXCLUDED.class_tag, |
| ratio_type = EXCLUDED.ratio_type, |
| matchup_key = EXCLUDED.matchup_key, |
| model_left = EXCLUDED.model_left, |
| model_right = EXCLUDED.model_right, |
| n_total = EXCLUDED.n_total, |
| n_humans = EXCLUDED.n_humans, |
| n_objects = EXCLUDED.n_objects, |
| prompt = EXCLUDED.prompt, |
| prompt_zh = EXCLUDED.prompt_zh, |
| humans = EXCLUDED.humans, |
| objects = EXCLUDED.objects |
| ''', |
| ( |
| c['id'], c['base_id'], c.get('seed_id', ''), int(c.get('level', 0)), c.get('class_tag', ''), |
| c.get('ratio_type', ''), c['matchup_key'], c['model_left'], c['model_right'], |
| int(c.get('n_total', 0)), int(c.get('n_humans', 0)), int(c.get('n_objects', 0)), |
| c.get('prompt', ''), c.get('prompt_zh', ''), json.dumps(c.get('humans', [])), json.dumps(c.get('objects', [])) |
| ), |
| ) |
|
|
|
|
| def upsert_assignments(assignments: list[dict]) -> None: |
| with get_conn() as conn: |
| with conn.cursor() as cur: |
| for a in assignments: |
| cur.execute( |
| f''' |
| INSERT INTO {TABLE_ASSIGNMENTS} |
| (combo_id, group_key, annotator_a_code, annotator_b_code) |
| VALUES (%s, %s, %s, %s) |
| ON CONFLICT (combo_id) DO UPDATE SET |
| group_key = EXCLUDED.group_key, |
| annotator_a_code = EXCLUDED.annotator_a_code, |
| annotator_b_code = EXCLUDED.annotator_b_code |
| ''', |
| (a['combo_id'], a['group_key'], int(a['annotator_a_code']), int(a['annotator_b_code'])), |
| ) |
|
|
