tarekziade HF Staff
Publish MIC usage, slim fixture and reviewed Transformers expectations
154c2d8 verified | # EmbeddingGemma: multimodal embeddings | |
| This page is the MIC integration contract for `google/embeddinggemma-2`. It shows the public `transformers` API for computing text and image embeddings and for running a short training step on the model's own parameters. The `runnable` blocks below are rendered by doc-builder and executed by the MIC runner against the slim fixture and the pinned checkpoint. Lines marked `# doc-builder: hide` are runner plumbing and are not shown in the rendered documentation. | |
| ## Compute text and image embeddings | |
| ```py runnable:embed | |
| import torch | |
| from transformers import AutoModel, AutoProcessor, set_seed | |
| import json # doc-builder: hide | |
| import os # doc-builder: hide | |
| model_id = "google/embeddinggemma-2" | |
| model_id = os.environ.get("MIC_MODEL", model_id) # doc-builder: hide | |
| variant = json.loads(os.environ.get("MIC_MODEL_KWARGS", "{}")) # doc-builder: hide | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| set_seed(0) # reproducible synthetic inputs and dropout | |
| processor = AutoProcessor.from_pretrained(model_id) | |
| model = AutoModel.from_pretrained(model_id).to(device) | |
| if variant: # doc-builder: hide | |
| del model # doc-builder: hide | |
| model = AutoModel.from_pretrained(model_id, **variant).to(device) # doc-builder: hide | |
| model.eval() | |
| texts = [ | |
| "What causes the northern lights?", | |
| "Charged particles from the sun cause the northern lights.", | |
| ] | |
| # Two small synthetic RGB images; the processor resizes and normalizes them. | |
| images = [torch.rand(3, 224, 224) for _ in range(2)] | |
| text_inputs = processor(text=texts, return_tensors="pt", padding=True) | |
| text_inputs = text_inputs.to(device) | |
| # The processor packs a list of images into one sequence: embed each on its own. | |
| image_inputs = [processor(images=[image], return_tensors="pt").to(device, dtype=model.dtype) for image in images] | |
| with torch.no_grad(): | |
| text_out = model(**text_inputs) | |
| image_outs = [model(**inputs) for inputs in image_inputs] | |
| # EmbeddingGemma2Model returns hidden states; mean-pool them into embeddings. | |
| def mean_pool(last_hidden_state, attention_mask): | |
| mask = attention_mask.unsqueeze(-1).to(last_hidden_state.dtype) | |
| return (last_hidden_state * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9) | |
| text_emb = mean_pool(text_out.last_hidden_state, text_inputs["attention_mask"]) | |
| image_emb = torch.cat([out.last_hidden_state.mean(dim=1) for out in image_outs]) # image tokens are never padded | |
| similarity = torch.nn.functional.cosine_similarity(text_emb[0], text_emb[1], dim=0) | |
| print(f"similarity: {similarity.item():.4f}") | |
| ``` | |
| ## One training step | |
| ```py runnable:embed:2 | |
| # In-batch contrastive step: each text should match its own image. The embeddings | |
| # are recomputed with gradients (the ones above were taken under no_grad). | |
| model.train() | |
| optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) | |
| text_train = mean_pool(model(**text_inputs).last_hidden_state, text_inputs["attention_mask"]) | |
| image_train = torch.cat([model(**inputs).last_hidden_state.mean(dim=1) for inputs in image_inputs]) | |
| scores = torch.nn.functional.normalize(text_train.float(), dim=-1) @ torch.nn.functional.normalize(image_train.float(), dim=-1).T | |
| loss = torch.nn.functional.cross_entropy(scores / 0.05, torch.arange(len(scores), device=device)) | |
| loss.backward() | |
| before = {n: p.detach().cpu().clone() for n, p in model.named_parameters() if p.requires_grad} # doc-builder: hide | |
| optimizer.step() | |
| print(f"loss: {loss.item():.4f}") | |
| if os.environ.get("MIC_RUNNER"): # doc-builder: hide | |
| named = dict(model.named_parameters()) | |
| checked = ( | |
| "language_model.layers.0.self_attn.q_proj.weight", | |
| "language_model.layers.0.mlp.down_proj.weight", | |
| "vision_tower.encoder.layers.0.self_attn.q_proj.linear.weight", | |
| ) | |
| result = { | |
| "runtime": { | |
| "device": model.device.type, | |
| "dtype": str(model.dtype).removeprefix("torch."), | |
| "attn_implementation": model.config._attn_implementation, | |
| }, | |
| "text_embeddings": text_emb.float().tolist(), | |
| "image_embeddings": image_emb.float().tolist(), | |
| "similarity": similarity.item(), | |
| "training": { | |
| "loss": loss.item(), | |
| "gradients": { | |
| n: bool( | |
| named[n].grad is not None | |
| and torch.count_nonzero(named[n].grad) > 0 | |
| and torch.isfinite(named[n].grad).all() | |
| ) | |
| for n in checked | |
| }, | |
| "updated": any( | |
| not torch.equal(p.detach().cpu(), before[n]) | |
| for n, p in model.named_parameters() | |
| if n in before | |
| ), | |
| }, | |
| } | |
| print(json.dumps(result)) | |
| ``` | |