Coverage for cafl4ds/filters/composite.py: 88%

1"""The composite selector — combine admission filters and a replay buffer (the ``A`` factor). 

2 

3The cheap Phase-0 knobs split into two composable roles (see :mod:`cafl4ds.filters.base`): 

4zero or more **admission** :class:`~cafl4ds.filters.base.Filter` s that narrow the incoming 

5batch, and at most one **replay** :class:`~cafl4ds.filters.base.ReplayBuffer` that stores and 

6replays. :class:`CompositeSelector` wires them behind the single ``Filter.select`` interface 

7the loop calls, so *what* combination of knobs is active is purely a config choice: 

8 

9====================== ========================= ==================== 

10Baseline admission buffer 

11====================== ========================= ==================== 

12B-floor — — 

13B1 (reservoir) — ReservoirReplay 

14dedup SemDeDup — 

15loss-gate (B3) LossGate — 

16B1.5 (dedup+replay) SemDeDup ReservoirReplay 

17====================== ========================= ==================== 

18 

19Admission filters run in listed order (each narrows what the next sees); the buffer always 

20runs last, on the admitted frames. This ordering is structural: dedup/loss-gate decide what is 

21*worth admitting*, and the buffer then decides what of the admitted history to *replay* — the 

22reverse (replaying, then deduping the replay) would be meaningless. 

23""" 

24 

25from __future__ import annotations 

26 

27from dataclasses import replace 

28 

29import torch 

30from loguru import logger 

31 

32from cafl4ds.data.streams import StreamBatch 

33from cafl4ds.filters.base import Filter, FilterContext, ReplayBuffer 

34 

35 

36class CompositeSelector(Filter): 

37 """Compose an admission-filter pipeline with an optional replay buffer.""" 

38 

39 def __init__(self, admission: list[Filter] | None = None, buffer: ReplayBuffer | None = None) -> None: 

40 """Build the composite. 

41 

42 Args: 

43 admission: Admission filters applied in order (each narrows the batch). ``None`` or 

44 an empty list means accept everything (no narrowing) — the B-floor admission. 

45 buffer: An optional replay buffer run after admission. ``None`` trains on the 

46 admitted frames directly (no replay). 

47 """ 

48 self.admission = list(admission or []) 

49 self.buffer = buffer 

50 

51 def select(self, batch: StreamBatch, ctx: FilterContext) -> torch.Tensor: 

52 """Run the admission pipeline, then the buffer, and return the training batch. 

53 

54 Args: 

55 batch: The incoming label-free stream batch. 

56 ctx: The live filter context (model, step), forwarded to every stage. 

57 

58 Returns: 

59 The images the SSL update trains on this step. May be *smaller* than the incoming 

60 batch (admission narrowed it, possibly to empty) or *larger* (the buffer replayed 

61 past frames). 

62 """ 

63 images = batch.images 

64 for stage in self.admission: 

65 images = stage.select(replace(batch, images=images), ctx) 

66 if images.shape[0] == 0: 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true

67 logger.debug(f"step {ctx.step}: admission emptied the batch at {type(stage).__name__}.") 

68 break 

69 if self.buffer is not None: 

70 images = self.buffer.observe(images, ctx) 

71 return images