Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-evaluation/src/lexigram/ai/evaluation/ablation.py: 50%

1"""Ablation runner for the AI evaluation framework. 

2 

3The :class:`AblationRunner` compares a baseline checkpoint against an 

4ablated one (a rerun with one configuration knob removed or changed) 

5and produces per-metric deltas plus a digest-stable result record. 

6""" 

7 

8from __future__ import annotations 

9 

10import hashlib 

11 

12from lexigram.ai.evaluation.exceptions import AblationError 

13from lexigram.ai.evaluation.tracking import canonical_json 

14from lexigram.contracts.ai.experiment import ( 

15 AblationResult, 

16 CheckpointStoreProtocol, 

17) 

18from lexigram.result import Err, Ok, Result 

19 

20_DIGEST_PREFIX = "ablation-" 

21 

22 

23class AblationRunner: 

24 """Compare baseline and ablated checkpoints from a checkpoint store. 

25 

26 Args: 

27 store: Checkpoint store holding the baseline and ablated payloads. 

28 """ 

29 

30 def __init__(self, store: CheckpointStoreProtocol) -> None: 

31 self._store = store 

32 

33 @staticmethod 

34 def deltas(before: dict[str, float], after: dict[str, float]) -> dict[str, float]: 

35 """Compute per-metric deltas (after minus before). 

36 

37 Args: 

38 before: Baseline metrics (name to value). 

39 after: Ablated metrics (name to value). 

40 

41 Returns: 

42 Metric name to delta mapping, including keys from either side. 

43 """ 

44 keys = list(before) + [key for key in after if key not in before] 

45 return {key: after.get(key, 0.0) - before.get(key, 0.0) for key in keys} 

46 

47 async def run( 

48 self, 

49 run_id: str, 

50 knob: str, 

51 baseline_slug: str, 

52 ablated_slug: str, 

53 ) -> Result[AblationResult, AblationError]: 

54 """Compare a baseline checkpoint against an ablated one in one run. 

55 

56 Args: 

57 run_id: Run the ablation was performed on. 

58 knob: The configuration knob that was ablated. 

59 baseline_slug: Checkpoint slug of the baseline run. 

60 ablated_slug: Checkpoint slug of the ablated run. 

61 

62 Returns: 

63 Ok(AblationResult) with per-metric deltas, or Err(AblationError) 

64 when either checkpoint is missing. 

65 """ 

66 return await self.compare(knob, run_id, baseline_slug, run_id, ablated_slug) 

67 

68 async def compare( 

69 self, 

70 knob: str, 

71 baseline_run_id: str, 

72 baseline_slug: str, 

73 ablated_run_id: str, 

74 ablated_slug: str, 

75 ) -> Result[AblationResult, AblationError]: 

76 """Compare checkpoints across two runs (e.g. control vs ablated). 

77 

78 Args: 

79 knob: The configuration knob that was ablated. 

80 baseline_run_id: Run identifier of the baseline. 

81 baseline_slug: Checkpoint slug of the baseline run. 

82 ablated_run_id: Run identifier of the ablated run. 

83 ablated_slug: Checkpoint slug of the ablated run. 

84 

85 Returns: 

86 Ok(AblationResult) with per-metric deltas, or Err(AblationError) 

87 when either checkpoint is missing. 

88 

89 Example: 

90 ```python 

91 runner = AblationRunner(store) 

92 result = await runner.compare( 

93 "thinking", 

94 "probe-42-a1b2c3d4", "baseline", 

95 "probe-42-9f8e7d6c", "ablated-thinking", 

96 ) 

97 ``` 

98 """ 

99 baseline = await self._store.load(baseline_run_id, baseline_slug) 

100 ablated = await self._store.load(ablated_run_id, ablated_slug) 

101 if baseline is None or ablated is None: 

102 missing = baseline_slug if baseline is None else ablated_slug 

103 return Err(AblationError(f"missing checkpoint {missing!r}")) 

104 before = {key: float(value) for key, value in baseline.payload.items()} 

105 after = {key: float(value) for key, value in ablated.payload.items()} 

106 deltas = self.deltas(before, after) 

107 digest = hashlib.sha256( 

108 canonical_json( 

109 { 

110 "knob": knob, 

111 "baseline_run_id": baseline_run_id, 

112 "baseline": baseline.digest, 

113 "ablated_run_id": ablated_run_id, 

114 "ablated": ablated.digest, 

115 "deltas": deltas, 

116 } 

117 ).encode() 

118 ).hexdigest() 

119 return Ok( 

120 AblationResult( 

121 run_id=baseline_run_id, 

122 knob=knob, 

123 baseline_slug=baseline_slug, 

124 ablated_slug=ablated_slug, 

125 deltas=deltas, 

126 digest=f"{_DIGEST_PREFIX}{digest}", 

127 ) 

128 ) 

129 

130 

131__all__ = ["AblationRunner"]