| """UCE Gemma: normalized embeddings or plain text to 80,513 labeled concepts.""" |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from safetensors.torch import load_file |
| from scipy import sparse |
|
|
| if __package__: |
| from .concept_metadata import ConceptMetadata, METADATA_FILES |
| else: |
| from concept_metadata import ConceptMetadata, METADATA_FILES |
|
|
|
|
| def _directory(path, revision=None, cache_dir=None, local_files_only=False, *, with_embedder=False): |
| directory = Path(path).expanduser() |
| if not directory.is_dir(): |
| from huggingface_hub import snapshot_download |
| directory = Path(snapshot_download( |
| repo_id=str(path), revision=revision, cache_dir=cache_dir, |
| local_files_only=local_files_only, |
| allow_patterns=METADATA_FILES + ['encoder.safetensors'] |
| + (['gemma_embedder/*'] if with_embedder else []), |
| )) |
| return directory |
|
|
|
|
| def _batch_size(value): |
| if not isinstance(value, int) or isinstance(value, bool) or value < 1: |
| raise ValueError('batch_size must be a positive integer') |
|
|
|
|
| class UCEGemmaEncoder(torch.nn.Module): |
| """Gemma embedding-to-concept encoder; float32 sigmoid scores, top 128.""" |
| def __init__(self, config, tensors, labels, *, prevalences=None, metadata=None): |
| super().__init__() |
| self.config = dict(config) |
| if config.get('model_type') != 'uce_gemma_80k' or config.get('format_version') != 1: |
| raise ValueError('Unsupported UCE Gemma package') |
| self.input_dim = config['input_dim'] |
| self.num_features = config['num_features'] |
| self.top_k = config['top_k'] |
| if config.get('input_unit_norm') is not True or not 0 < self.top_k <= self.num_features: |
| raise ValueError('Invalid encoder configuration') |
| if set(tensors) != {'weight', 'bias'}: |
| raise ValueError('Expected weight and bias only') |
| if tensors['weight'].shape != (self.num_features, self.input_dim) or tensors['bias'].shape != (self.num_features,): |
| raise ValueError('Invalid encoder weight shape') |
| for name, tensor in tensors.items(): |
| if tensor.dtype != torch.float32 or not torch.isfinite(tensor).all(): |
| raise ValueError('Expected finite float32 weights') |
| self.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False)) |
| self.labels = dict(labels) |
| self.metadata = metadata or ConceptMetadata(config, labels, prevalences=prevalences) |
| self._prevalences = self.metadata._prevalences |
| if set(self.labels) != set(range(self.num_features)) or any(not label.strip() for label in self.labels.values()): |
| raise ValueError('Every output coordinate must have one nonempty label') |
| self.eval() |
|
|
| @classmethod |
| def from_pretrained(cls, model_id_or_path, *, revision=None, device='cpu', cache_dir=None, local_files_only=False): |
| directory = _directory(model_id_or_path, revision, cache_dir, local_files_only) |
| config = json.loads((directory / 'config.json').read_text()) |
| metadata = ConceptMetadata.from_pretrained(directory) |
| if not Path(model_id_or_path).expanduser().is_dir(): |
| metadata.provenance.update( |
| source_repo_id=str(model_id_or_path), requested_revision=revision, |
| resolved_revision=directory.name if directory.parent.name == 'snapshots' else None) |
| return cls(config, load_file(str(directory / 'encoder.safetensors')), metadata.labels, |
| metadata=metadata).to(device) |
|
|
| def get_prevalence(self, concept_id): |
| """Return saved historical Gemini baseline counts, fractions, and bounds.""" |
| return self.metadata.get_prevalence(concept_id) |
|
|
| @torch.inference_mode() |
| def forward(self, embeddings): |
| """Return device tensors (concept IDs, scores), each [batch, 128].""" |
| if self.weight.dtype != torch.float32 or self.bias.dtype != torch.float32: |
| raise ValueError('Keep the concept encoder in float32') |
| if isinstance(embeddings, torch.Tensor): |
| if embeddings.is_complex() or embeddings.dtype == torch.bool: |
| raise ValueError('Expected real embeddings') |
| x = embeddings.detach().to(device=self.weight.device, dtype=torch.float32) |
| else: |
| a = np.asarray(embeddings) |
| if a.dtype.kind not in 'fiu': |
| raise ValueError('Expected real embeddings') |
| x = torch.from_numpy(np.array(a, dtype=np.float32, copy=True)).to(self.weight.device) |
| if x.ndim != 2 or x.shape[1] != self.input_dim or not torch.isfinite(x).all(): |
| raise ValueError(f'Expected finite embeddings of shape (N, {self.input_dim})') |
| if not torch.allclose(torch.linalg.vector_norm(x, dim=1), torch.ones(len(x), device=x.device), atol=2e-5, rtol=0): |
| raise ValueError('Provide unit-normalized Gemma document embeddings') |
| with torch.autocast(device_type=self.weight.device.type, enabled=False): |
| logits = torch.nn.functional.linear(x, self.weight, self.bias) |
| if not torch.isfinite(logits).all(): |
| raise ValueError('Nonfinite encoder scores') |
| |
| values, indices = logits.topk(self.top_k, dim=-1) |
| return indices, torch.sigmoid(values) |
|
|
| def encode_iter(self, embeddings, *, batch_size=64): |
| """Yield float32 SciPy CSR batches with 80,513 concept columns.""" |
| _batch_size(batch_size) |
| if not isinstance(embeddings, (np.ndarray, torch.Tensor)): |
| raise TypeError('Expected a NumPy array or Torch tensor') |
| if embeddings.ndim != 2 or embeddings.shape[1] != self.input_dim: |
| raise ValueError(f'Expected shape (N, {self.input_dim})') |
| for start in range(0, len(embeddings), batch_size): |
| ids, vals = self(embeddings[start:start + batch_size]) |
| ids, vals = ids.cpu().numpy().astype(np.int32), vals.cpu().numpy() |
| rowptr = np.arange(0, vals.size + 1, self.top_k, dtype=np.int32) |
| result = sparse.csr_matrix((vals.ravel(), ids.ravel(), rowptr), shape=(len(vals), self.num_features)) |
| result.eliminate_zeros() |
| result.sort_indices() |
| yield result |
|
|
| def encode(self, embeddings, *, batch_size=64): |
| batches = list(self.encode_iter(embeddings, batch_size=batch_size)) |
| return sparse.vstack(batches, format='csr') if batches else sparse.csr_matrix((0, self.num_features), dtype=np.float32) |
|
|
| def explain(self, row, *, top_n=20, include_prevalence=False): |
| """Describe an encoded row using saved labels; no LLM/API calls.""" |
| if not sparse.issparse(row) or row.shape != (1, self.num_features): |
| raise ValueError('Pass one encoded sparse row') |
| if not isinstance(top_n, int) or isinstance(top_n, bool) or top_n < 0: |
| raise ValueError('top_n must be a nonnegative integer') |
| row = row.tocsr(copy=True) |
| row.sum_duplicates(); row.eliminate_zeros() |
| order = np.lexsort((row.indices, -row.data))[:top_n] |
| concepts = [dict(concept_id=int(row.indices[i]), label=self.labels[int(row.indices[i])], activation=float(row.data[i])) for i in order] |
| if include_prevalence: |
| for concept in concepts: |
| concept['baseline_prevalence'] = self.get_prevalence(concept['concept_id']) |
| return concepts |
|
|
|
|
| class UCEGemma(torch.nn.Module): |
| """Plain strings -> text-only Gemma backbone -> concept scores. |
| |
| Includes no vision/audio towers, multimodal projection weights, or processor. |
| Uses the original bfloat16 text inference and float32 concept encoder. |
| """ |
| def __init__(self, encoder, text_model, tokenizer): |
| super().__init__() |
| self.encoder = encoder |
| self.text_model = text_model |
| self.tokenizer = tokenizer |
| self.text_config = encoder.config['text_pipeline'] |
| self.config = encoder.config |
| self.num_features = encoder.num_features |
| self.labels = encoder.labels |
| self.metadata = encoder.metadata |
| self.eval() |
|
|
| @classmethod |
| def from_pretrained(cls, model_id_or_path, *, revision=None, device='cpu', cache_dir=None, local_files_only=False): |
| from transformers import AutoTokenizer, EmbeddingGemma2TextModel |
| directory = _directory(model_id_or_path, revision, cache_dir, local_files_only, with_embedder=True) |
| encoder = UCEGemmaEncoder.from_pretrained(directory, device=device) |
| |
| if not Path(model_id_or_path).expanduser().is_dir(): |
| encoder.metadata.provenance.update( |
| source_repo_id=str(model_id_or_path), requested_revision=revision, |
| resolved_revision=directory.name if directory.parent.name == 'snapshots' else None) |
| if 'text_pipeline' not in encoder.config or not (directory / 'gemma_embedder').is_dir(): |
| raise ValueError('Load the full package to encode text') |
| text_model = EmbeddingGemma2TextModel.from_pretrained( |
| directory / 'gemma_embedder', dtype=torch.bfloat16, local_files_only=True, |
| ).to(device).eval() |
| tokenizer = AutoTokenizer.from_pretrained(directory / 'gemma_embedder', local_files_only=True, trust_remote_code=False) |
| return cls(encoder, text_model, tokenizer) |
|
|
| @staticmethod |
| def _texts(texts): |
| if isinstance(texts, str): |
| texts = [texts] |
| if not isinstance(texts, (list, tuple)) or any(not isinstance(t, str) for t in texts): |
| raise TypeError('Expected a string or a list/tuple of plain strings') |
| return texts |
|
|
| @torch.inference_mode() |
| def embed(self, texts, *, batch_size=32): |
| """Return normalized float32 [N,768] embeddings; overlength inputs raise.""" |
| texts = self._texts(texts) |
| _batch_size(batch_size) |
| batches = [] |
| device = next(self.text_model.parameters()).device |
| for start in range(0, len(texts), batch_size): |
| inputs = self.tokenizer([self.text_config['prompt'] + t for t in texts[start:start+batch_size]], |
| padding=True, truncation=False, return_tensors='pt') |
| if inputs['input_ids'].shape[1] > self.text_config['max_seq_length']: |
| raise ValueError('Text exceeds 8192 tokens including prompt; split the document explicitly') |
| ids, mask = inputs['input_ids'].to(device), inputs['attention_mask'].to(device) |
| |
| model_ids = ids.clone() |
| for token_id in self.text_config['placeholder_token_ids']: |
| model_ids.masked_fill_(ids == token_id, self.tokenizer.pad_token_id) |
| hidden = self.text_model(input_ids=model_ids, attention_mask=mask).last_hidden_state |
| expanded = mask.unsqueeze(-1).expand_as(hidden).to(hidden.dtype) |
| pooled = (hidden * expanded).sum(1) / expanded.sum(1).clamp_min(1e-9) |
| |
| pooled = torch.nn.functional.normalize(pooled, p=2, dim=1) |
| pooled = torch.nn.functional.normalize(pooled, p=2, dim=1) |
| array = pooled.float().cpu().numpy() |
| array /= np.linalg.norm(array, axis=1, keepdims=True) |
| batches.append(array) |
| return np.concatenate(batches) if batches else np.empty((0, self.encoder.input_dim), np.float32) |
|
|
| def encode_iter(self, texts, *, batch_size=32): |
| texts = self._texts(texts) |
| _batch_size(batch_size) |
| for start in range(0, len(texts), batch_size): |
| embeddings = self.embed(texts[start:start+batch_size], batch_size=batch_size) |
| yield self.encoder.encode(embeddings, batch_size=batch_size) |
|
|
| def encode(self, texts, *, batch_size=32): |
| batches = list(self.encode_iter(texts, batch_size=batch_size)) |
| return sparse.vstack(batches, format='csr') if batches else sparse.csr_matrix((0, self.encoder.num_features), dtype=np.float32) |
|
|
| def forward(self, texts): |
| return self.encode(texts) |
|
|
| def explain(self, row, *, top_n=20, include_prevalence=False): |
| return self.encoder.explain(row, top_n=top_n, include_prevalence=include_prevalence) |
|
|
| def get_prevalence(self, concept_id): |
| """Look up one concept's historical Gemini corpus baseline statistics.""" |
| return self.encoder.get_prevalence(concept_id) |
|
|
| def main(): |
| import argparse |
| parser = argparse.ArgumentParser(description='Universal Concept Encoder: Gemma text to 80,513 concepts') |
| parser.add_argument('--model', default=str(Path(__file__).resolve().parent), help='Local package or Hub repository ID') |
| parser.add_argument('--text', required=True, help='Plain text to encode') |
| parser.add_argument('--device', default='cpu') |
| parser.add_argument('--top-n', type=int, default=20) |
| args = parser.parse_args() |
| model = UCEGemma.from_pretrained(args.model, device=args.device) |
| print(json.dumps(model.explain(model.encode(args.text)[0], top_n=args.top_n), ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
