Coverage for transformer_lens/tools/model_registry/verification.py: 98%

1"""Verification tracking for model compatibility. 

2 

3This module provides dataclasses and utilities for tracking which models 

4have been verified to work with TransformerLens. 

5""" 

6 

7from dataclasses import dataclass, field 

8from datetime import date, datetime 

9from typing import Optional 

10 

11 

12@dataclass 

13class VerificationRecord: 

14 """A record of a model verification. 

15 

16 Attributes: 

17 model_id: The HuggingFace model ID that was verified 

18 architecture_id: The architecture type of the model 

19 verified_date: Date when verification was performed 

20 verified_by: Who performed the verification (user, CI, etc.) 

21 transformerlens_version: Version of TransformerLens used 

22 notes: Optional notes about the verification 

23 invalidated: Whether this verification has been invalidated 

24 invalidation_reason: Reason for invalidation if applicable 

25 """ 

26 

27 model_id: str 

28 verified_date: date 

29 architecture_id: str = "Unknown" 

30 verified_by: Optional[str] = None 

31 transformerlens_version: Optional[str] = None 

32 # P4 verdict flips are undiagnosable without knowing which profile and 

33 # scoring scale produced the record. 

34 prompt_profile: Optional[str] = None 

35 p4_scoring_version: Optional[int] = None 

36 notes: Optional[str] = None 

37 invalidated: bool = False 

38 invalidation_reason: Optional[str] = None 

39 

40 def to_dict(self) -> dict: 

41 """Convert to a JSON-serializable dictionary.""" 

42 return { 

43 "model_id": self.model_id, 

44 "architecture_id": self.architecture_id, 

45 "verified_date": self.verified_date.isoformat(), 

46 "verified_by": self.verified_by, 

47 "transformerlens_version": self.transformerlens_version, 

48 "prompt_profile": self.prompt_profile, 

49 "p4_scoring_version": self.p4_scoring_version, 

50 "notes": self.notes, 

51 "invalidated": self.invalidated, 

52 "invalidation_reason": self.invalidation_reason, 

53 } 

54 

55 @classmethod 

56 def from_dict(cls, data: dict) -> "VerificationRecord": 

57 """Create from a dictionary.""" 

58 return cls( 

59 model_id=data["model_id"], 

60 architecture_id=data.get("architecture_id", "Unknown"), 

61 verified_date=date.fromisoformat(data["verified_date"]), 

62 verified_by=data.get("verified_by"), 

63 transformerlens_version=data.get("transformerlens_version"), 

64 prompt_profile=data.get("prompt_profile"), 

65 p4_scoring_version=data.get("p4_scoring_version"), 

66 notes=data.get("notes"), 

67 invalidated=data.get("invalidated", False), 

68 invalidation_reason=data.get("invalidation_reason"), 

69 ) 

70 

71 

72@dataclass 

73class VerificationHistory: 

74 """History of all model verifications. 

75 

76 Attributes: 

77 records: List of all verification records 

78 last_updated: When this history was last updated 

79 """ 

80 

81 records: list[VerificationRecord] = field(default_factory=list) 

82 last_updated: Optional[datetime] = None 

83 

84 def to_dict(self) -> dict: 

85 """Convert to a JSON-serializable dictionary.""" 

86 return { 

87 "last_updated": self.last_updated.isoformat() if self.last_updated else None, 

88 "records": [r.to_dict() for r in self.records], 

89 } 

90 

91 @classmethod 

92 def from_dict(cls, data: dict) -> "VerificationHistory": 

93 """Create from a dictionary.""" 

94 last_updated = None 

95 if data.get("last_updated"): 95 ↛ 97line 95 didn't jump to line 97 because the condition on line 95 was always true

96 last_updated = datetime.fromisoformat(data["last_updated"]) 

97 return cls( 

98 records=[VerificationRecord.from_dict(r) for r in data.get("records", [])], 

99 last_updated=last_updated, 

100 ) 

101 

102 def get_record(self, model_id: str) -> Optional[VerificationRecord]: 

103 """Get the most recent valid verification record for a model. 

104 

105 Args: 

106 model_id: The model ID to look up 

107 

108 Returns: 

109 The verification record, or None if not found or invalidated 

110 """ 

111 for record in reversed(self.records): 

112 if record.model_id == model_id and not record.invalidated: 

113 return record 

114 return None 

115 

116 def is_verified(self, model_id: str) -> bool: 

117 """Check if a model has a valid verification. 

118 

119 Args: 

120 model_id: The model ID to check 

121 

122 Returns: 

123 True if the model has a valid (non-invalidated) verification 

124 """ 

125 return self.get_record(model_id) is not None 

126 

127 def add_record(self, record: VerificationRecord) -> None: 

128 """Add a new verification record. 

129 

130 Args: 

131 record: The verification record to add 

132 """ 

133 self.records.append(record) 

134 self.last_updated = datetime.now() 

135 

136 def invalidate(self, model_id: str, reason: str) -> bool: 

137 """Invalidate the most recent verification for a model. 

138 

139 Args: 

140 model_id: The model ID to invalidate 

141 reason: Reason for invalidation 

142 

143 Returns: 

144 True if a record was invalidated, False if not found 

145 """ 

146 record = self.get_record(model_id) 

147 if record: 

148 record.invalidated = True 

149 record.invalidation_reason = reason 

150 self.last_updated = datetime.now() 

151 return True 

152 return False