generate_tasks.py

from jinja2 import Environment, FileSystemLoader, select_autoescape
from bson import json_util
import uuid
import json
import random
import os
import re


def get_html_for_pair(itemA, itemB):
    env = Environment(
        loader=FileSystemLoader(os.path.dirname(os.path.abspath(__file__))),
        autoescape=select_autoescape(enabled_extensions=("html", "j2"))
    )
    tpl = env.get_template("pair_panels.html.j2")
    return tpl.render(a=itemA, b=itemB)


def slugify_name(name):
    return re.sub(r"[^a-z0-9]+", "_", name.lower()).strip("_")


def load_enriched_pairs(path, sample_size=None):
    with open(path, "r", encoding="utf-8") as fh:
        all_pairs = json.load(fh)

    if sample_size is None or sample_size >= len(all_pairs):
        return all_pairs

    selected_ids = random.sample(list(all_pairs.keys()), sample_size)
    return {k: all_pairs[k] for k in selected_ids}


def assign_pairs_to_annotators(pairs, annotators):
    """
    Assign each pair to exactly one annotator, balancing workload.

    Returns:
        tasks_by_annotator: {annotator_name: [Label Studio tasks]}
        tracking_records: flat list with one row per pair assignment
    """
    tasks_by_annotator = {annotator: [] for annotator in annotators}
    task_counts = {annotator: 0 for annotator in annotators}
    tracking_records = []

    pair_items = list(pairs.items())
    random.shuffle(pair_items)

    for pair_uid, (rawA, rawB) in pair_items:
        annotator = min(
            annotators,
            key=lambda a: (task_counts[a], random.random())
        )

        html = get_html_for_pair(rawA, rawB)
        task_id = str(uuid.uuid4())

        task = {
            "id": task_id,
            "data": {
                "pair_uid": pair_uid,
                "conflict_id": pair_uid.split("__pair_")[0],
                "annotator": annotator,
                "html": html
            }
        }

        tasks_by_annotator[annotator].append(task)

        tracking_records.append({
            "pair_uid": pair_uid,
            "conflict_id": pair_uid.split("__pair_")[0],
            "itemA": rawA,
            "itemB": rawB,
            "html": html,
            "annotator": annotator,
            "task_id": task_id
        })

        task_counts[annotator] += 1

    return tasks_by_annotator, tracking_records


def write_split_task_files(tasks_by_annotator, output_dir, chunk_size=30):
    """
    Write one or more JSON files per annotator, split into chunks of `chunk_size`.
    """
    os.makedirs(output_dir, exist_ok=True)

    for annotator, tasks in tasks_by_annotator.items():
        annotator_slug = slugify_name(annotator)

        for chunk_idx, start in enumerate(range(0, len(tasks), chunk_size), start=1):
            chunk = tasks[start:start + chunk_size]

            filename = f"tasks_{annotator_slug}_part_{chunk_idx:02d}.json"
            path = os.path.join(output_dir, filename)

            with open(path, "w", encoding="utf-8") as f:
                json.dump(chunk, f, ensure_ascii=False, indent=2)

            print(
                f"{annotator}: part {chunk_idx:02d} "
                f"({len(chunk)} tasks) -> {path}"
            )


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate Label Studio annotation tasks from pre-built enriched pairs."
    )
    parser.add_argument(
        "--pairs",
        required=True,
        help="Path to the enriched-pairs JSON file produced by build_enriched_pairs.py.",
    )
    parser.add_argument(
        "--sample-size",
        type=int,
        default=None,
        help="Number of pairs to sample for annotation. If omitted, all pairs are used.",
    )
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=30,
        help="Number of tasks per output file.",
    )
    parser.add_argument(
        "--round",
        default="round_1",
        help="Annotation round name. Tasks and manifest are written under rounds/<round>/.",
    )

    args = parser.parse_args()

    annotators = [
        "ARNOLD AMUSENGERI",
        "AINHOA DE LA TORRE MONTANA",
        "ANDREA MORALES MATA",
        "ANI LISETH VALLE BANEGAS",
        "EVA MARTIN DEL PICO",
        "JUAN CRISANTOS PEREZ CARRASCO",
        "LAURA ANTIGA",
        "LAURA PORTELL SILVA",
        "MARIA MORALES MARTINEZ",
        "PAULA IBORRA DE TOLEDO",
        "PABLO IGNACIO MARTINEZ DIAZ",
        "SILVIA DI GIORGIO",
        "AINA JENNE",
        "SALVADOR CAPELLA GUTIÉRREZ",
        "MARIA PAOLA FERRI",
        "DIMITRY REPCHEVSKY"
    ]

    random.seed(42)

    pairs = load_enriched_pairs(args.pairs, sample_size=args.sample_size)

    tasks_by_annotator, tracking_records = assign_pairs_to_annotators(
        pairs,
        annotators
    )

    round_dir = os.path.join("rounds", args.round)
    output_dir = os.path.join(round_dir, "tasks")
    os.makedirs(round_dir, exist_ok=True)

    write_split_task_files(
        tasks_by_annotator,
        output_dir=output_dir,
        chunk_size=args.chunk_size
    )

    manifest_path = os.path.join(round_dir, "manifest.json")
    with open(manifest_path, "w", encoding="utf-8") as f:
        json.dump(
            tracking_records,
            f,
            ensure_ascii=False,
            indent=2,
            default=json_util.default
        )

    print(f"Tracking manifest written with {len(tracking_records)} pair assignments.")