Blablablab/audio-classification

1"""2Iterative Best-Worst Scaling (IBWS) Manager3 4Implements the IBWS algorithm from "Baby Bear: Seeking a Just Right Rating Scale5for Scalar Annotations" (arxiv 2408.09765). IBWS extends standard BWS with a6Quicksort-like adaptive loop:7 81. Round 1: Generate tuples from full pool, annotators select best/worst92. Score items, partition into upper/middle/lower buckets103. Round N: Generate tuples WITHIN each bucket, annotate, partition again114. Stop when all buckets are terminal (< tuple_size items) or max_rounds reached12 13Output: Ordinal ranking from bucket positions + within-bucket scores.14 15Usage:16    from potato.ibws_manager import get_ibws_manager, init_ibws_manager17 18    mgr = init_ibws_manager(config, pool_items, id_key, text_key)19    round1_tuples = mgr.get_current_round_tuples()20 21    # After annotations are complete for current round:22    if mgr.check_round_complete(ism, bws_schema_name):23        new_tuples = mgr.advance_round(ism, bws_schema_name)24        # Add new_tuples to ISM25"""26 27import logging28import math29import threading30from typing import Any, Dict, List, Optional, Tuple31 32from potato.bws_scoring import BwsScorer33from potato.bws_tuple_generator import BwsTupleGenerator34 35logger = logging.getLogger(__name__)36 37# Singleton instance38_ibws_manager = None39_ibws_lock = threading.Lock()40 41 42def init_ibws_manager(config: dict, pool_items: List[Dict[str, Any]],43                      id_key: str, text_key: str) -> "IBWSManager":44    """Initialize the global IBWS manager singleton."""45    global _ibws_manager46    with _ibws_lock:47        _ibws_manager = IBWSManager(config, pool_items, id_key, text_key)48    return _ibws_manager49 50 51def get_ibws_manager() -> Optional["IBWSManager"]:52    """Get the global IBWS manager (None if not initialized)."""53    return _ibws_manager54 55 56def clear_ibws_manager():57    """Clear the global IBWS manager (for testing)."""58    global _ibws_manager59    with _ibws_lock:60        _ibws_manager = None61 62 63class IBWSManager:64    """Manages iterative BWS rounds, partitioning, and tuple generation."""65 66    def __init__(self, config: dict, pool_items: List[Dict[str, Any]],67                 id_key: str, text_key: str):68        self._lock = threading.RLock()69 70        self.id_key = id_key71        self.text_key = text_key72 73        ibws_config = config["ibws_config"]74        self.tuple_size = ibws_config.get("tuple_size", 4)75        self.max_rounds = ibws_config.get("max_rounds", None)  # None = auto76        self.seed = ibws_config.get("seed", 42)77        self.scoring_method = ibws_config.get("scoring_method", "counting")78        self.tuples_per_item_per_round = ibws_config.get("tuples_per_item_per_round", 2)79 80        # Store original pool items with their IDs81        self.pool_items = list(pool_items)82        self.pool_item_map = {str(item[id_key]): item for item in pool_items}83 84        # Partition state: list of buckets, each bucket is a list of item IDs85        # Start with one bucket containing all items86        self.current_round = 0  # 0 = not started, 1 = round 1 active, etc.87        self.buckets: List[List[str]] = [[str(item[id_key]) for item in pool_items]]88        self.terminal_buckets: List[List[str]] = []  # Buckets too small to partition further89 90        # Track tuples for each round: round_num -> list of tuple IDs91        self.round_tuples: Dict[int, List[str]] = {}92 93        # Track all generated tuples' data for scoring94        self.tuple_data: Dict[str, Dict[str, Any]] = {}95 96        # Completed flag97        self.completed = False98 99    def generate_round_tuples(self) -> List[Dict[str, Any]]:100        """Generate tuples for the next round from current active buckets.101 102        Returns list of tuple instance dicts ready for ISM.add_item().103        """104        with self._lock:105            self.current_round += 1106            round_num = self.current_round107 108            all_tuples = []109            tuple_ids = []110 111            new_buckets = []112            for bucket_idx, bucket_item_ids in enumerate(self.buckets):113                if len(bucket_item_ids) < self.tuple_size:114                    # Terminal bucket — too few items to form a tuple115                    self.terminal_buckets.append(bucket_item_ids)116                    continue117 118                new_buckets.append(bucket_item_ids)119 120                # Build pool items for this bucket121                bucket_pool = [self.pool_item_map[iid] for iid in bucket_item_ids122                               if iid in self.pool_item_map]123 124                if len(bucket_pool) < self.tuple_size:125                    self.terminal_buckets.append(bucket_item_ids)126                    continue127 128                # Calculate tuples needed for this bucket129                min_appearances = self.tuples_per_item_per_round * self.tuple_size130                num_tuples = max(1, math.ceil(131                    len(bucket_pool) * self.tuples_per_item_per_round / self.tuple_size132                ))133 134                prefix = f"ibws_r{round_num}_b{bucket_idx}"135                generator = BwsTupleGenerator(136                    pool_items=bucket_pool,137                    id_key=self.id_key,138                    text_key=self.text_key,139                    tuple_size=self.tuple_size,140                    num_tuples=num_tuples,141                    seed=self.seed + round_num * 1000 + bucket_idx,142                    min_item_appearances=min_appearances,143                )144 145                tuples = generator.generate()146 147                # Rename tuple IDs with our prefix148                for i, t in enumerate(tuples):149                    new_id = f"{prefix}_{i + 1:04d}"150                    t[self.id_key] = new_id151                    t["_ibws_round"] = round_num152                    t["_ibws_bucket"] = bucket_idx153                    self.tuple_data[new_id] = t154                    tuple_ids.append(new_id)155 156                all_tuples.extend(tuples)157 158            # Update active buckets (excluding those that became terminal)159            self.buckets = new_buckets160            self.round_tuples[round_num] = tuple_ids161 162            if not all_tuples:163                # All buckets are terminal164                self.completed = True165 166            logger.info(167                f"IBWS round {round_num}: Generated {len(all_tuples)} tuples "168                f"across {len(self.buckets)} active buckets "169                f"({len(self.terminal_buckets)} terminal)"170            )171 172            return all_tuples173 174    def check_round_complete(self, ism, bws_schema_name: str) -> bool:175        """Check if all tuples in the current round have been annotated.176 177        Uses ISM's instance_annotators tracking to see if each tuple178        has at least one annotator.179 180        Args:181            ism: ItemStateManager instance182            bws_schema_name: Name of the BWS annotation schema183 184        Returns:185            True if all current round tuples have at least one annotation186        """187        with self._lock:188            if self.completed or self.current_round == 0:189                return False190 191            round_tuple_ids = self.round_tuples.get(self.current_round, [])192            if not round_tuple_ids:193                return False194 195            # Check that every tuple in this round has at least one annotator196            for tuple_id in round_tuple_ids:197                annotators = ism.instance_annotators.get(tuple_id, set())198                if not annotators:199                    return False200 201            return True202 203    def advance_round(self, ism, usm, bws_schema_name: str) -> List[Dict[str, Any]]:204        """Score current round, partition buckets, generate next round tuples.205 206        Args:207            ism: ItemStateManager instance208            usm: UserStateManager instance209            bws_schema_name: Name of the BWS annotation schema210 211        Returns:212            List of new tuple instance dicts for the next round (empty if done)213        """214        with self._lock:215            if self.completed:216                return []217 218            if self.max_rounds and self.current_round >= self.max_rounds:219                self.completed = True220                logger.info(f"IBWS: Reached max_rounds ({self.max_rounds}), stopping")221                return []222 223            # Score current round and partition each active bucket224            new_buckets = []225            for bucket_idx, bucket_item_ids in enumerate(self.buckets):226                if len(bucket_item_ids) < self.tuple_size:227                    self.terminal_buckets.append(bucket_item_ids)228                    continue229 230                # Collect annotations for tuples that contain items from this bucket231                annotations = self._collect_bucket_annotations(232                    bucket_item_ids, ism, usm, bws_schema_name233                )234 235                if not annotations:236                    # No annotations — can't partition, keep bucket as-is237                    new_buckets.append(bucket_item_ids)238                    continue239 240                # Score items in this bucket241                bucket_pool = [self.pool_item_map[iid] for iid in bucket_item_ids242                               if iid in self.pool_item_map]243                scorer = BwsScorer(annotations, bucket_pool, self.id_key, self.text_key)244                scores = scorer.score(self.scoring_method)245 246                # Partition into upper/middle/lower thirds247                upper, middle, lower = self._partition_bucket(bucket_item_ids, scores)248 249                for sub_bucket in [upper, middle, lower]:250                    if sub_bucket:251                        new_buckets.append(sub_bucket)252 253            self.buckets = new_buckets254 255            # Check if all remaining buckets are terminal256            active_count = sum(1 for b in self.buckets if len(b) >= self.tuple_size)257            if active_count == 0:258                # Move remaining small buckets to terminal259                for b in self.buckets:260                    if len(b) < self.tuple_size:261                        self.terminal_buckets.append(b)262                self.buckets = []263                self.completed = True264                logger.info("IBWS: All buckets terminal, annotation complete")265                return []266 267            # Generate tuples for the next round268            return self.generate_round_tuples()269 270    def _collect_bucket_annotations(self, bucket_item_ids, ism, usm,271                                     bws_schema_name: str) -> List[Dict[str, Any]]:272        """Collect BWS annotations for tuples containing items from a bucket."""273        bucket_id_set = set(bucket_item_ids)274        annotations = []275 276        # Look at tuples from the current round277        round_tuple_ids = self.round_tuples.get(self.current_round, [])278 279        for tuple_id in round_tuple_ids:280            tuple_info = self.tuple_data.get(tuple_id)281            if not tuple_info:282                continue283 284            bws_items = tuple_info.get("_bws_items", [])285            # Check if this tuple's items overlap with our bucket286            tuple_source_ids = {item["source_id"] for item in bws_items}287            if not tuple_source_ids.intersection(bucket_id_set):288                continue289 290            # Collect annotations from all users for this tuple291            for user_state in usm.get_all_users():292                username = user_state.get_user_id()293                label_store = getattr(user_state, 'instance_id_to_label_to_value', {})294 295                if tuple_id not in label_store:296                    continue297 298                labels = label_store[tuple_id]299                best_val = None300                worst_val = None301                for label_obj, value in labels.items():302                    if label_obj.get_schema() == bws_schema_name:303                        if label_obj.get_name() == "best":304                            best_val = value305                        elif label_obj.get_name() == "worst":306                            worst_val = value307 308                if best_val and worst_val:309                    annotations.append({310                        "instance_id": tuple_id,311                        "bws_items": bws_items,312                        "best": best_val,313                        "worst": worst_val,314                        "annotator": username,315                    })316 317        return annotations318 319    def _partition_bucket(self, item_ids: List[str],320                          scores: Dict[str, Dict[str, Any]]) -> Tuple[List[str], List[str], List[str]]:321        """Partition a bucket into upper/middle/lower thirds by score.322 323        Uses equal-thirds of sorted list (not score thresholds) for balanced partitions.324        """325        # Sort by score descending326        sorted_ids = sorted(327            item_ids,328            key=lambda iid: scores.get(iid, {}).get("score", 0.0),329            reverse=True330        )331 332        n = len(sorted_ids)333        third = n // 3334 335        # Handle remainder: distribute extra items to middle336        upper = sorted_ids[:third]337        lower = sorted_ids[n - third:] if third > 0 else []338        middle = sorted_ids[third:n - third] if third > 0 else sorted_ids339 340        return upper, middle, lower341 342    def get_round_info(self) -> Dict[str, Any]:343        """Get current round information for UI display."""344        with self._lock:345            total_tuples_this_round = len(self.round_tuples.get(self.current_round, []))346            active_buckets = len([b for b in self.buckets if len(b) >= self.tuple_size])347            terminal_count = len(self.terminal_buckets)348            total_items = len(self.pool_items)349 350            # Items in terminal buckets (already ranked)351            terminal_items = sum(len(b) for b in self.terminal_buckets)352 353            return {354                "current_round": self.current_round,355                "max_rounds": self.max_rounds,356                "total_tuples_this_round": total_tuples_this_round,357                "active_buckets": active_buckets,358                "terminal_buckets": terminal_count,359                "total_items": total_items,360                "terminal_items": terminal_items,361                "completed": self.completed,362            }363 364    def get_final_ranking(self) -> List[Dict[str, Any]]:365        """Produce final ordinal ranking from bucket positions + within-bucket scores.366 367        Returns list of dicts sorted by rank:368            [{"item_id": str, "rank": int, "bucket_position": int, "text": str}, ...]369        """370        with self._lock:371            # Combine terminal buckets (ordered by when they became terminal = higher quality)372            # and any remaining active buckets373            all_buckets = list(self.terminal_buckets) + list(self.buckets)374 375            ranking = []376            rank = 1377            for bucket_position, bucket in enumerate(all_buckets):378                for item_id in bucket:379                    item = self.pool_item_map.get(item_id, {})380                    ranking.append({381                        "item_id": item_id,382                        "rank": rank,383                        "bucket_position": bucket_position,384                        "text": str(item.get(self.text_key, "")),385                    })386                    rank += 1387 388            return ranking389 390    def is_completed(self) -> bool:391        """Check if IBWS has completed all rounds."""392        with self._lock:393            return self.completed394