Coverage for imcluster/llm.py: 100.00%

1"""LLM-assisted names for image clusters.""" 

2 

3from collections.abc import Sequence 

4from typing import Any 

5 

6import llmloader 

7import numpy as np 

8from langchain_core.messages import HumanMessage, SystemMessage 

9from numpy.typing import ArrayLike 

10from rich.console import Console 

11from rich.markup import escape 

12 

13from .io import ImclusterIO 

14 

15console = Console() 

16 

17DEFAULT_LLM = "gpt-5.6-luna" 

18 

19 

20def encode_message(text: str) -> dict[str, str]: 

21 """Return a multimodal text-content block.""" 

22 return {"type": "text", "text": text} 

23 

24 

25def encode_thumbnail_message(thumbnail: str) -> dict[str, Any]: 

26 """Return a multimodal image block for a cached JPEG thumbnail.""" 

27 return { 

28 "type": "image_url", 

29 "image_url": {"url": f"data:image/jpeg;base64,{thumbnail}"}, 

30 } 

31 

32 

33def name_cluster( 

34 llm: Any, 

35 in_group: Sequence[str], 

36 out_group: Sequence[str] | None = None, 

37) -> str: 

38 """Generate a concise name from cached thumbnails in and outside a cluster. 

39 

40 Args: 

41 llm: LangChain-compatible multimodal chat model. 

42 in_group: JPEG thumbnails belonging to the cluster, encoded as base64. 

43 out_group: Optional thumbnails from nearby images outside the cluster. 

44 

45 Returns: 

46 A stripped, non-empty cluster name. 

47 

48 Raises: 

49 ValueError: If no in-cluster thumbnails or no name is returned. 

50 """ 

51 if not in_group: 

52 raise ValueError("At least one in-cluster thumbnail is required") 

53 if out_group: 

54 opening = ( 

55 "Name the visual category shared by the images in this cluster. " 

56 "Images outside the cluster are contrasting examples and should help " 

57 "you distinguish the category." 

58 ) 

59 group_heading = "Images inside the cluster:" 

60 else: 

61 opening = "Name the visual category shared by these images." 

62 group_heading = "Images:" 

63 content: list[str | dict[Any, Any]] = [ 

64 encode_message(opening), 

65 encode_message(group_heading), 

66 ] 

67 content.extend(encode_thumbnail_message(thumbnail) for thumbnail in in_group) 

68 if out_group: 

69 content.append(encode_message("Nearby images outside the cluster:")) 

70 content.extend(encode_thumbnail_message(thumbnail) for thumbnail in out_group) 

71 content.append( 

72 encode_message( 

73 "Respond with only a short descriptive cluster name. Do not add quotes, " 

74 "a prefix, a sentence, or punctuation." 

75 ) 

76 ) 

77 response = llm.invoke( 

78 [ 

79 SystemMessage( 

80 content=("You assign concise, specific names to clusters of images.") 

81 ), 

82 HumanMessage(content=content), 

83 ] 

84 ) 

85 response_content = getattr(response, "content", response) 

86 if isinstance(response_content, list): 

87 response_content = "".join( 

88 str(block.get("text", "")) if isinstance(block, dict) else str(block) 

89 for block in response_content 

90 ) 

91 name = str(response_content).strip().strip("\"'").strip() 

92 if not name: 

93 raise ValueError("The LLM returned an empty cluster name") 

94 return name 

95 

96 

97def _normalize_vectors(feature_vectors: ArrayLike, image_count: int) -> Any: 

98 """Validate and cosine-normalize image feature vectors.""" 

99 vectors = np.asarray(feature_vectors, dtype=float) 

100 if vectors.ndim != 2 or len(vectors) != image_count: 

101 raise ValueError("feature_vectors must contain one row per image") 

102 norms = np.linalg.norm(vectors, axis=1, keepdims=True) 

103 return np.divide(vectors, norms, out=np.zeros_like(vectors), where=norms != 0) 

104 

105 

106def _representative_positions( 

107 positions: Any, 

108 normalized: Any, 

109 limit: int, 

110) -> list[int]: 

111 """Select a medoid followed by diverse images using farthest-first traversal.""" 

112 cluster_vectors = normalized[positions] 

113 medoid_offset = int(np.argmax(cluster_vectors @ cluster_vectors.sum(axis=0))) 

114 selected = [int(positions[medoid_offset])] 

115 while len(selected) < min(limit, len(positions)): 

116 similarities = normalized[positions] @ normalized[selected].T 

117 nearest_selected = similarities.max(axis=1) 

118 for selected_position in selected: 

119 nearest_selected[positions == selected_position] = np.inf 

120 selected.append(int(positions[np.argmin(nearest_selected)])) 

121 return selected 

122 

123 

124def _outside_positions( 

125 cluster_positions: Any, 

126 representatives: Sequence[int], 

127 normalized: Any, 

128 limit: int, 

129) -> list[int]: 

130 """Select outside images nearest to any representative image.""" 

131 if limit == 0: 

132 return [] 

133 outside = np.setdiff1d( 

134 np.arange(len(normalized)), cluster_positions, assume_unique=True 

135 ) 

136 if not len(outside): 

137 return [] 

138 similarities = normalized[outside] @ normalized[list(representatives)].T 

139 scores = similarities.max(axis=1) 

140 ranked = np.argsort(-scores, kind="stable")[:limit] 

141 return [int(position) for position in outside[ranked]] 

142 

143 

144def name_clusters( 

145 imcluster_io: ImclusterIO, 

146 feature_vectors: ArrayLike, 

147 cluster_column: str, 

148 llm: Any = DEFAULT_LLM, 

149 temperature: float = 0.2, 

150 api_key: str | None = None, 

151 in_group_size: int = 10, 

152 out_group_size: int = 0, 

153 force: bool = False, 

154) -> dict[object, str]: 

155 """Name clusters from representative cached thumbnails and save the results. 

156 

157 Args: 

158 imcluster_io: Image cache containing cluster labels and thumbnails. 

159 feature_vectors: Original image embeddings used to choose prompt examples. 

160 cluster_column: DataFrame column containing cluster assignments. 

161 llm: Loaded multimodal model or an llmloader model identifier. 

162 temperature: Sampling temperature used when loading a model identifier. 

163 api_key: Optional provider API key passed to llmloader. 

164 in_group_size: Maximum representative thumbnails from inside each cluster. 

165 out_group_size: Maximum nearby contrasting thumbnails outside each cluster. 

166 force: Regenerate names even when a complete name cache exists. 

167 

168 Returns: 

169 Mapping from cluster labels to their generated display names. 

170 

171 Raises: 

172 ValueError: If required columns, vectors, thumbnails, or sizes are invalid. 

173 """ 

174 if cluster_column not in imcluster_io.df: 

175 raise ValueError(f"Missing clustering results column: {cluster_column}") 

176 if "thumbnail" not in imcluster_io.df: 

177 raise ValueError("Cached thumbnails are required to name clusters") 

178 if in_group_size < 1: 

179 raise ValueError("in_group_size must be at least 1") 

180 if out_group_size < 0: 

181 raise ValueError("out_group_size must not be negative") 

182 

183 name_column = f"{cluster_column}_name" 

184 labels = imcluster_io.df[cluster_column].to_numpy() 

185 unique_labels = list(dict.fromkeys(labels.tolist())) 

186 if imcluster_io.has_column(name_column) and not force: 

187 cached_names: dict[object, str] = {} 

188 complete = True 

189 for label in unique_labels: 

190 values = { 

191 value.strip() 

192 for value in imcluster_io.df.loc[labels == label, name_column].tolist() 

193 if isinstance(value, str) and value.strip() 

194 } 

195 if len(values) != 1: 

196 complete = False 

197 break 

198 cached_names[label] = next(iter(values)) 

199 if complete: 

200 console.print( 

201 f"[green]Using cached cluster names:[/green] loaded " 

202 f"{len(cached_names)} names from '{imcluster_io.output}'." 

203 ) 

204 return cached_names 

205 

206 normalized = _normalize_vectors(feature_vectors, len(imcluster_io.images)) 

207 thumbnails = imcluster_io.df["thumbnail"].tolist() 

208 if any(not isinstance(thumbnail, str) or not thumbnail for thumbnail in thumbnails): 

209 raise ValueError("Every image must have a cached thumbnail before naming") 

210 if isinstance(llm, str): 

211 load_kwargs: dict[str, Any] = {"temperature": temperature} 

212 if api_key is not None: 

213 load_kwargs["api_key"] = api_key 

214 llm = llmloader.load(llm, **load_kwargs) 

215 

216 names: dict[object, str] = {} 

217 output_names = np.empty(len(labels), dtype=object) 

218 with console.status("[cyan]Generating descriptive cluster names...[/cyan]"): 

219 for label in unique_labels: 

220 cluster_positions = np.flatnonzero(labels == label) 

221 if label == -1: 

222 cluster_name = "Noise" 

223 else: 

224 representatives = _representative_positions( 

225 cluster_positions, normalized, in_group_size 

226 ) 

227 outside = _outside_positions( 

228 cluster_positions, representatives, normalized, out_group_size 

229 ) 

230 cluster_name = name_cluster( 

231 llm, 

232 [thumbnails[position] for position in representatives], 

233 [thumbnails[position] for position in outside], 

234 ) 

235 if label == -1: 

236 display_label = "noise" 

237 elif isinstance(label, (int, np.integer)): 

238 display_label = f"cluster {int(label) + 1}" 

239 else: 

240 display_label = f"cluster {label}" 

241 console.print( 

242 f"[green]Named {escape(display_label)}:[/green] {escape(cluster_name)}" 

243 ) 

244 names[label] = cluster_name 

245 output_names[cluster_positions] = cluster_name 

246 imcluster_io.save_column(name_column, output_names.tolist()) 

247 console.print( 

248 f"[green]Wrote cluster names:[/green] saved {len(names)} names to " 

249 f"'{imcluster_io.output}'." 

250 ) 

251 return names