Blablablab/audio-classification

1"""2Best-Worst Scaling Score Estimation3 4Computes item scores from BWS annotations using three methods:51. Counting: score = (best_count - worst_count) / appearances  (no dependencies)62. Bradley-Terry: pairwise comparison model via choix  (requires choix)73. Plackett-Luce: partial ranking model via choix  (requires choix)8 9Usage as library:10    from potato.bws_scoring import BwsScorer11    scorer = BwsScorer(annotations, pool_items, id_key)12    scores = scorer.counting()13 14Usage as CLI:15    python -m potato.bws_scoring --config config.yaml --method counting16"""17 18import argparse19import csv20import json21import logging22import os23import sys24from typing import Any, Dict, List, Optional, Tuple25 26logger = logging.getLogger(__name__)27 28 29class BwsScorer:30    """Compute BWS scores from annotations."""31 32    def __init__(33        self,34        annotations: List[Dict[str, Any]],35        pool_items: List[Dict[str, Any]],36        id_key: str,37        text_key: str = "text",38    ):39        """40        Args:41            annotations: List of annotation dicts, each with:42                - "instance_id": tuple instance ID (e.g. "bws_tuple_0001")43                - "bws_items": list of {source_id, text, position}44                - "best": position label (e.g. "B")45                - "worst": position label (e.g. "D")46                - "annotator": username47            pool_items: Original pool items48            id_key: Key for item IDs in pool_items49            text_key: Key for item text in pool_items50        """51        self.annotations = annotations52        self.pool_items = pool_items53        self.id_key = id_key54        self.text_key = text_key55 56        # Build item index57        self.item_ids = [str(item[id_key]) for item in pool_items]58        self.item_texts = {59            str(item[id_key]): str(item.get(text_key, ""))60            for item in pool_items61        }62        self.item_id_to_idx = {iid: idx for idx, iid in enumerate(self.item_ids)}63 64    def _resolve_annotation(65        self, ann: Dict[str, Any]66    ) -> Optional[Tuple[str, str, List[str]]]:67        """Resolve an annotation to (best_source_id, worst_source_id, all_source_ids).68 69        Returns None if annotation is incomplete.70        """71        best_pos = ann.get("best")72        worst_pos = ann.get("worst")73        bws_items = ann.get("bws_items", [])74 75        if not best_pos or not worst_pos or not bws_items:76            return None77 78        pos_to_id = {item["position"]: item["source_id"] for item in bws_items}79        best_id = pos_to_id.get(best_pos)80        worst_id = pos_to_id.get(worst_pos)81 82        if not best_id or not worst_id:83            return None84 85        all_ids = [item["source_id"] for item in bws_items]86        return best_id, worst_id, all_ids87 88    def counting(self) -> Dict[str, Dict[str, Any]]:89        """Counting method: score = (best_count - worst_count) / appearances.90 91        Returns dict mapping item_id to {score, best_count, worst_count, appearances, text}.92        """93        best_counts = {iid: 0 for iid in self.item_ids}94        worst_counts = {iid: 0 for iid in self.item_ids}95        appearances = {iid: 0 for iid in self.item_ids}96 97        for ann in self.annotations:98            resolved = self._resolve_annotation(ann)99            if not resolved:100                continue101 102            best_id, worst_id, all_ids = resolved103            for iid in all_ids:104                if iid in appearances:105                    appearances[iid] += 1106            if best_id in best_counts:107                best_counts[best_id] += 1108            if worst_id in worst_counts:109                worst_counts[worst_id] += 1110 111        scores = {}112        for iid in self.item_ids:113            app = appearances[iid]114            if app > 0:115                score = (best_counts[iid] - worst_counts[iid]) / app116            else:117                score = 0.0118 119            scores[iid] = {120                "score": score,121                "best_count": best_counts[iid],122                "worst_count": worst_counts[iid],123                "appearances": app,124                "text": self.item_texts.get(iid, ""),125            }126 127        return scores128 129    def bradley_terry(self) -> Dict[str, Dict[str, Any]]:130        """Bradley-Terry model via choix.131 132        Converts each BWS annotation to pairwise comparisons:133        - Best item beats every other item (K-1 comparisons)134        - Every item beats the worst item (K-1 comparisons)135        """136        try:137            import choix138        except ImportError:139            raise ImportError(140                "Bradley-Terry scoring requires the 'choix' package. "141                "Install it with: pip install choix"142            )143 144        n_items = len(self.item_ids)145        comparisons = []146 147        for ann in self.annotations:148            resolved = self._resolve_annotation(ann)149            if not resolved:150                continue151 152            best_id, worst_id, all_ids = resolved153            best_idx = self.item_id_to_idx.get(best_id)154            worst_idx = self.item_id_to_idx.get(worst_id)155 156            if best_idx is None or worst_idx is None:157                continue158 159            # Best beats all others160            for iid in all_ids:161                idx = self.item_id_to_idx.get(iid)162                if idx is not None and idx != best_idx:163                    comparisons.append((best_idx, idx))164 165            # All others beat worst166            for iid in all_ids:167                idx = self.item_id_to_idx.get(iid)168                if idx is not None and idx != worst_idx:169                    comparisons.append((idx, worst_idx))170 171        if not comparisons:172            return {173                iid: {"score": 0.0, "text": self.item_texts.get(iid, "")}174                for iid in self.item_ids175            }176 177        params = choix.ilsr_pairwise(n_items, comparisons, alpha=0.01)178 179        scores = {}180        for iid in self.item_ids:181            idx = self.item_id_to_idx[iid]182            scores[iid] = {183                "score": float(params[idx]),184                "text": self.item_texts.get(iid, ""),185            }186 187        return scores188 189    def plackett_luce(self) -> Dict[str, Dict[str, Any]]:190        """Plackett-Luce model via choix.191 192        Converts BWS to partial rankings:193        Each annotation yields top-1 (best) selections, processed via ilsr_top1.194        """195        try:196            import choix197        except ImportError:198            raise ImportError(199                "Plackett-Luce scoring requires the 'choix' package. "200                "Install it with: pip install choix"201            )202 203        n_items = len(self.item_ids)204        # Use pairwise comparisons to approximate partial rankings205        # Best > middle items, middle items > worst206        comparisons = []207 208        for ann in self.annotations:209            resolved = self._resolve_annotation(ann)210            if not resolved:211                continue212 213            best_id, worst_id, all_ids = resolved214            best_idx = self.item_id_to_idx.get(best_id)215            worst_idx = self.item_id_to_idx.get(worst_id)216 217            if best_idx is None or worst_idx is None:218                continue219 220            middle_ids = [221                iid for iid in all_ids if iid != best_id and iid != worst_id222            ]223 224            # Best beats all middle items225            for iid in middle_ids:226                idx = self.item_id_to_idx.get(iid)227                if idx is not None:228                    comparisons.append((best_idx, idx))229 230            # All middle items beat worst231            for iid in middle_ids:232                idx = self.item_id_to_idx.get(iid)233                if idx is not None:234                    comparisons.append((idx, worst_idx))235 236            # Best beats worst237            comparisons.append((best_idx, worst_idx))238 239        if not comparisons:240            return {241                iid: {"score": 0.0, "text": self.item_texts.get(iid, "")}242                for iid in self.item_ids243            }244 245        params = choix.ilsr_pairwise(n_items, comparisons, alpha=0.01)246 247        scores = {}248        for iid in self.item_ids:249            idx = self.item_id_to_idx[iid]250            scores[iid] = {251                "score": float(params[idx]),252                "text": self.item_texts.get(iid, ""),253            }254 255        return scores256 257    def score(self, method: str = "counting") -> Dict[str, Dict[str, Any]]:258        """Compute scores using the specified method."""259        if method == "counting":260            return self.counting()261        elif method == "bradley_terry":262            return self.bradley_terry()263        elif method == "plackett_luce":264            return self.plackett_luce()265        else:266            raise ValueError(267                f"Unknown scoring method: {method}. "268                "Use 'counting', 'bradley_terry', or 'plackett_luce'."269            )270 271 272def write_scores(273    scores: Dict[str, Dict[str, Any]],274    output_path: str,275) -> None:276    """Write scores to a TSV file.277 278    Output columns: item_id, text, score, best_count, worst_count, appearances, rank279    """280    # Sort by score descending281    sorted_items = sorted(scores.items(), key=lambda x: x[1]["score"], reverse=True)282 283    os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)284 285    with open(output_path, "w", newline="") as f:286        writer = csv.writer(f, delimiter="\t")287        writer.writerow(288            ["item_id", "text", "score", "best_count", "worst_count", "appearances", "rank"]289        )290        for rank, (item_id, data) in enumerate(sorted_items, 1):291            writer.writerow([292                item_id,293                data.get("text", ""),294                f"{data['score']:.6f}",295                data.get("best_count", ""),296                data.get("worst_count", ""),297                data.get("appearances", ""),298                rank,299            ])300 301    logger.info(f"Wrote BWS scores to {output_path}")302 303 304def collect_annotations_from_output(305    output_dir: str, bws_schema_name: str, config: dict306) -> List[Dict[str, Any]]:307    """Collect BWS annotations from Potato's output directory.308 309    Reads annotation files and reconstructs BWS annotation records.310    """311    annotations = []312    pool_items_by_tuple = {}313 314    # Get pool items from config315    bws_pool = config.get("_bws_pool_items", [])316    id_key = config["item_properties"]["id_key"]317 318    # We need to read the saved annotations from the output dir319    # Potato saves annotations as {output_dir}/{annotator}.jsonl320    if not os.path.isdir(output_dir):321        logger.warning(f"Output directory not found: {output_dir}")322        return annotations323 324    for fname in os.listdir(output_dir):325        if not fname.endswith(".jsonl"):326            continue327 328        annotator = fname.replace(".jsonl", "")329        fpath = os.path.join(output_dir, fname)330 331        with open(fpath, "r") as f:332            for line in f:333                line = line.strip()334                if not line:335                    continue336                try:337                    record = json.loads(line)338                except json.JSONDecodeError:339                    continue340 341                instance_id = record.get("id")342                ann_data = record.get("annotation", {})343 344                # Look for BWS schema annotations345                best_val = None346                worst_val = None347                for schema_name, schema_ann in ann_data.items():348                    if schema_name == bws_schema_name:349                        best_val = schema_ann.get("best")350                        worst_val = schema_ann.get("worst")351                        break352 353                if not best_val or not worst_val:354                    continue355 356                # Get BWS items from the instance data357                bws_items = record.get("_bws_items", [])358 359                annotations.append({360                    "instance_id": instance_id,361                    "bws_items": bws_items,362                    "best": best_val,363                    "worst": worst_val,364                    "annotator": annotator,365                })366 367    return annotations368 369 370def main():371    """CLI entry point for BWS scoring."""372    parser = argparse.ArgumentParser(373        description="Compute BWS scores from Potato annotation output"374    )375    parser.add_argument(376        "--config", required=True, help="Path to Potato config YAML file"377    )378    parser.add_argument(379        "--method",380        default="counting",381        choices=["counting", "bradley_terry", "plackett_luce"],382        help="Scoring method (default: counting)",383    )384    parser.add_argument(385        "--output",386        default=None,387        help="Output TSV file path (default: {output_dir}/bws_scores.tsv)",388    )389    args = parser.parse_args()390 391    logging.basicConfig(level=logging.INFO)392 393    # Load config394    import yaml395 396    with open(args.config, "r") as f:397        config = yaml.safe_load(f)398 399    output_dir = config.get("output_annotation_dir", "annotation_output")400    id_key = config["item_properties"]["id_key"]401    text_key = config["item_properties"]["text_key"]402 403    # Find BWS schema name404    bws_schema_name = None405    for scheme in config.get("annotation_schemes", []):406        if scheme.get("annotation_type") == "bws":407            bws_schema_name = scheme["name"]408            break409 410    if not bws_schema_name:411        print("Error: No BWS annotation scheme found in config", file=sys.stderr)412        sys.exit(1)413 414    # Load pool items from data files415    pool_items = []416    for data_file in config.get("data_files", []):417        if isinstance(data_file, dict):418            data_file = data_file.get("path")419        if not data_file:420            continue421 422        with open(data_file, "r") as f:423            if data_file.endswith(".json"):424                pool_items.extend(json.load(f))425            else:426                for line in f:427                    line = line.strip()428                    if line:429                        pool_items.append(json.loads(line))430 431    # Collect annotations432    annotations = collect_annotations_from_output(output_dir, bws_schema_name, config)433 434    if not annotations:435        print("No BWS annotations found in output directory", file=sys.stderr)436        sys.exit(1)437 438    print(f"Found {len(annotations)} BWS annotations for {len(pool_items)} pool items")439 440    # Score441    scorer = BwsScorer(annotations, pool_items, id_key, text_key)442    scores = scorer.score(args.method)443 444    # Write output445    output_path = args.output or os.path.join(output_dir, "bws_scores.tsv")446    write_scores(scores, output_path)447    print(f"Scores written to {output_path}")448 449 450if __name__ == "__main__":451    main()452