klpeng/uce-gemma-80k

UCE Gemma

Convert text into scores for 80,513 named concepts.

Install

Choose the dependencies for your workflow. This repository provides standalone Python modules; it is not installed by pip install uce-gemma.

For metadata only (including DataFrame exports):

pip install numpy huggingface_hub pandas
hf download klpeng/uce-gemma-80k uce_gemma.py concept_metadata.py --local-dir uce-gemma-metadata
cd uce-gemma-metadata

Then run the metadata-only example below. ConceptMetadata.from_pretrained() downloads the metadata files automatically without downloading model weights. For reproducible analyses, use the same pinned Hub commit for the runtime files and the revision argument to from_pretrained().

For the full text encoder:

pip install -U huggingface_hub
hf download klpeng/uce-gemma-80k --local-dir uce-gemma-80k
cd uce-gemma-80k
pip install -r requirements.txt

Run

Run Python from the downloaded directory:

from uce_gemma import UCEGemma

model = UCEGemma.from_pretrained(".", device="cuda")
scores = model.encode("Djokovic overwhelms Nadal for Wimbledon Title")

for concept in model.explain(scores[0], top_n=20):
    print(concept["concept_id"], concept["label"], concept["activation"])

Use device="cpu" without a CUDA GPU. Pass a list of strings to encode multiple texts. Inputs can contain up to 8,192 tokens, including the document prompt.

scores is a sparse matrix of shape (number_of_texts, 80513), retaining the 128 highest-scoring concepts per text. Values are sigmoid scores, not calibrated probabilities.

Baseline prevalences

prevalences.csv and prevalences.npz contain the same baseline statistics for all 80,513 concepts, joined to labels.csv and output columns by concept_id. They include counts, prevalence fractions, and one-sided 99% Clopper-Pearson upper bounds at activation thresholds >0.05 and >0.1.

These files are copied unchanged from Gemini UCD. They describe 89,827,558 stored Gemini corpus rows, including duplicates. They were computed using the historical 131,072-feature Gemini encoder's global top-128 selection, then restricted and remapped to these 80,513 concept IDs. They are historical Gemini teacher baselines, not prevalences measured from UCE Gemma outputs; the thresholds apply to the historical Gemini activations, not Gemma sigmoid scores. Use them as heuristic reference rates for concept enrichment.

prevalences.json records the original population, computation, and confidence bound assumptions. package_provenance.json records the source repository, pinned revision, and file hashes.

# Look up one concept's historical Gemini corpus baseline.
baseline = model.get_prevalence(0)
print(baseline["label"])
print(baseline["prevalence_gt_0p05"])          # fraction, not percent
print(baseline["prevalence_gt_0p05_upper_99"])

# Include baseline statistics alongside each explained concept.
for concept in model.explain(scores[0], top_n=20, include_prevalence=True):
    baseline = concept["baseline_prevalence"]
    print(concept["label"], concept["activation"], baseline["prevalence_gt_0p05"])

get_prevalence(concept_id) returns a dictionary containing concept_id, label, n_corpus_rows, and the count_gt_*, prevalence_gt_*, and prevalence_gt_*_upper_99 fields for thresholds 0p05 and 0p1. Both methods work on UCEGemma and UCEGemmaEncoder. Remote loading downloads baseline data automatically. Existing explain calls keep their original output unless include_prevalence=True is passed.

Metadata access

Full encoding workflow:

from uce_gemma import UCEGemma

model = UCEGemma.from_pretrained("klpeng/uce-gemma-80k")
X = model.encode(texts)
metadata = model.metadata

Metadata-only workflow:

from uce_gemma import ConceptMetadata

metadata = ConceptMetadata.from_pretrained("klpeng/uce-gemma-80k")
metadata.search("democracy")
metadata.describe([42, 135])

Metadata loading downloads only labels, saved statistics, config, and provenance; it imports neither PyTorch nor Transformers and does not load weights. Install numpy and huggingface_hub for remote metadata access, and pandas for table exports. For a standalone download, also obtain the small runtime files uce_gemma.py and concept_metadata.py and run Python from that directory. Full encoding additionally needs _model.py and the encoder dependencies.

label(id) returns a label; search(text) returns a list of concept_id/label dictionaries using case-insensitive substring matching. get_prevalence(id) and get_fidelity(id) return dictionaries of original saved statistics. describe(ids) returns a pandas DataFrame in requested order; to_dataframe() returns all concepts in activation-column order. Both tables contain concept_id, label, original prevalence fields, and fidelity fields prefixed with fidelity_. Undefined fidelity metrics are None in dictionaries and missing values in tables. fidelity_feature_id is a historical source ID; always join activations using concept_id.

metadata.available_thresholds  # (0.05, 0.1) for the current release
metadata.prevalence(42, threshold=0.05)  # saved prevalence fraction
metadata.prevalence(42, threshold=0.05, bound="upper", confidence=0.99)
metadata.get_fidelity(42)
metadata.to_dataframe().to_csv("concepts.csv", index=False)
metadata.save_pretrained("./uce-metadata")
restored = ConceptMetadata.from_pretrained("./uce-metadata")

Thresholds apply to historical Gemini teacher activations. Bounds keep their original statistical definitions; unavailable thresholds, bounds, or confidence levels raise errors. Omitting confidence for a bound uses the config's saved confidence level. No bounds are recomputed. Packages without optional statistics still support labels and tables; individual statistic getters raise FileNotFoundError when that metadata is absent.

Both loaders accept local directories, revision, cache_dir, and local_files_only. Pin revision when matching a previously encoded matrix. metadata.provenance records the label SHA-256 and, for Hub loads, the repository, requested revision, and resolved snapshot revision when available. Metadata exports preserve the original config (including model provenance), labels, statistic files, and method descriptions, plus metadata_provenance.json. Existing model.explain() and model.get_prevalence() keep their return formats; UCEGemmaEncoder also exposes the same metadata interface.

Per-concept fidelities

fidelities.csv contains precision, recall, F1, accuracy, positive/zero YES rates, and confusion counts for all 80,513 concepts. Join to labels.csv or encoder output columns using concept_id; feature_id is a historical source coordinate. fidelity_summary.json contains aggregate results, and fidelities.json records the evaluation method and source hashes.

Reference: Gemma raw sigmoid >0.92278 versus absence from stored sparse top 128; up to 30 positives and 30 sparse zeros per concept. GPT-6 Luna with reasoning disabled judged held-out test texts (256-word cap). Annotation YES is the prediction; activating versus zero is the reference. These conditional positive/zero sample metrics are automated judgments and are not weighted by corpus prevalence. Empty recall/F1 cells indicate concepts with no positive examples; these are excluded from corresponding averages.

import csv

with open("fidelities.csv", newline="") as f:
    fidelities = {int(r["concept_id"]): r for r in csv.DictReader(f)}
value = fidelities[0]["f1"]
f1 = float(value) if value else None