File size: 7,931 Bytes
e5242f5
d5de844
 
 
 
 
 
 
 
 
e5242f5
d5de844
 
 
ed04ebd
d5de844
 
 
 
 
 
 
 
 
5edc52b
d5de844
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5edc52b
d5de844
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
---
library_name: transformers
pipeline_tag: feature-extraction
tags:
  - metaencoder
  - muse_glimmer
  - multimodal-embedding
  - image-text-retrieval
  - video-retrieval
  - document-retrieval
---

# MetaEncoder v1

![](benchmarks.png)

## Usage

Requires `transformers>=5.15` (the first release with `muse_glimmer`); see `requirements.txt`. One 80 GB GPU holds
the model in bf16.

```python
from modeling_metaencoder import MetaEncoder

model = MetaEncoder.from_pretrained("facebook/meta-encoder-v1")
```

`MetaEncoder` has three calls. Callers pass content and task descriptions; the wrapper writes every prompt itself, in
the layout the encoder was trained on (see [Prompt formatting](#prompt-formatting)).

| Call | Use it to | Returns |
| :--- | :--- | :--- |
| `encode(inputs, instruction=None)` | embed inputs to store, cluster or compare | `(n, 6656)` float32 tensor, L2-normalised |
| `retrieve(query, corpus, top_k=None, query_instruction=None, corpus_instruction=None)` | rank a corpus against a query | `list[Hit]`, best first |
| `decide(state, questions, images=None, video=None)` | answer Jev-style questions about a state | Jev answers with probabilities |

### Inputs

`encode` and `retrieve` take *items*. An item is a string, a PIL image, or a dict with any of `text`, `image`,
`images`, `video` (a list of frames) and `instruction`. The instruction says what the embedding is for, e.g.
"Represent the given image." or "Find a document image that matches the given query."; the same input embeds
differently under different instructions, and queries and documents usually take different ones.

### 1. Encode

```python
from PIL import Image

docs = model.encode([{"image": Image.open("chart.png")}, {"image": Image.open("slide.png")}],
                    instruction="Understand the content of the provided document image.")
MetaEncoder.similarity(docs, docs)   # cosine similarity matrix
```

`inputs` is one item or a list of them; `instruction` applies to the items that carry none of their own.

### 2. Retrieve

```python
hits = model.retrieve("Which quarter had the highest revenue?", docs, top_k=1,
                      query_instruction="Find a document image that matches the given query.")
hits[0].index, hits[0].score, hits[0].item   # position in the corpus, cosine similarity, the corpus item
```

`corpus` is a list of items, or the tensor `encode` returned for them: encode a large corpus once and search the
tensor repeatedly.

### 3. Decide

`decide` takes a Jev decision request, a state and named questions, and answers each question:

```python
result = model.decide(
    {"query": "used iphone 15"},
    {
        "shopping": {"type": "noul", "instructions": "Is the user shopping for a product?"},
        "topic": {"type": "choice", "instructions": "Which topic is the query about?",
                  "criteria": {"electronics": "", "food": "", "travel": "", "other": ""}},
        "intent": {"type": "score", "instructions": "How strong is the purchase intent?",
                   "criteria": ["none", "weak", "moderate", "strong"]},
    },
)
result["questions"]["shopping"]   # {"value": True, "probabilities": {"false": 0.016, "true": 0.984},
                                  #  "confidence": 0.881, "scores": {...}}
result["questions"]["topic"]["value"], result["questions"]["intent"]["value"]   # "electronics", 1.89
```

- `state`: what the questions are about. A string, or any JSON value (an object of fields, a list of records).
- `images`, `video`: the state's media: a PIL image or a list of them, and a list of video frames.
- `questions`: `{name: {"type", "instructions", "criteria"}}`, with `instructions` the question itself:

| `type` | `criteria` | `value` | `probabilities` keys |
| :--- | :--- | :--- | :--- |
| `noul` | optional `{"true": ..., "false": ...}` descriptions | `True` / `False` | `"false"`, `"true"` |
| `choice` | `{option: description}`, 2 to 26 options; descriptions may be empty | the chosen option | the options |
| `score` | the ordered list of level labels, e.g. `["none", "weak", "strong"]` | expected level index (float) | `"0"` .. `"n-1"` |

Each answer also has `confidence` (1 minus the normalised entropy of `probabilities`) and `scores` (the raw cosine
similarities). Questions are answered independently; for a question that only applies given another's answer,
decide the first and ignore the second as needed.

`decide_batch(states, questions, images=None, video=None)` answers the same questions for many states (one
`images` / `video` entry per state).

Probabilities are `softmax(score / 0.02)`, the contrastive training temperature. To fit the temperature to your
data, decide a labelled set and call `model.calibrate(answers, labels)` with the per-question answers
(`result["questions"][name]`) and their correct keys (the option, `True`/`False`, or the level index); later
`decide` calls use the fitted temperature.


## Prompt formatting

The wrapper formats every input as in training, so callers never write prompts:

- **System turn** (`chat_template.jinja`): the fixed header ending in "Represent the user's input.", then the item's
  instruction after a blank line when it has one. The date line is fixed, so an input embeds the same on any day.
- **User turn**: images and video frames first, then the text. Images keep their native resolution, snapped to a
  28-pixel grid, up to 4096 merged tokens. Video is encoded frame by frame; more than 8 frames are subsampled
  uniformly to 8, as in training.
- **Pooling**: the last position, the user turn's `<|eot|>`, with left padding. Inputs are never truncated (that
  would drop the pooled token); training sequences were at most 16,384 tokens.

`decide` turns each question into a closed-set decision row of training, scored as retrieval over its options:

| Part | Rendered as |
| :--- | :--- |
| query text | the state, a blank line, `Task: {instructions}`, then the options: `Options: a, b, c`, or one `- option` per line when an option is long or punctuated |
| state | a string as is; a JSON value with `json.dumps` |
| query instruction | "Read the task definition and select the option it calls for." for a string state; "Answer the question about the scenario by selecting the correct option." for a JSON state |
| options | `noul`: `No, Yes`; `choice`: the criteria keys; `score`: the level labels |
| candidates | one per option: the option, then its description (if any) on the next line |
| candidate instruction | `Option for the task: {instructions}` when every option is a short label (at most 5 words), else none |

Option descriptions are optional. The yes/no, choice and score rows of training used bare option labels; candidate
descriptions were trained only as records of the entities that options name.

### Without the wrapper

`modeling_metaencoder.py` only formats the prompt and pools; the same embedding with public classes directly:

```python
import torch, torch.nn.functional as F
from safetensors.torch import load_file
from transformers import AutoModel, AutoProcessor

model = AutoModel.from_pretrained(repo, dtype=torch.bfloat16, device_map="auto").eval()   # MuseGlimmerModel
processor = AutoProcessor.from_pretrained(repo)
processor.tokenizer.padding_side = "left"
head = load_file(f"{repo}/embedding_head.safetensors")["embedding_head.weight"].to(model.device, torch.bfloat16)

messages = [{"role": "system", "content": "Represent the given image."},          # optional instruction
            {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "a dog"}]}]
text = processor.apply_chat_template(messages, tokenize=False)
inputs = processor(text=[text], images=[image], add_special_tokens=False, max_image_tokens=4096,
                   return_tensors="pt").to(model.device)
with torch.no_grad():
    hidden = model(**inputs).last_hidden_state[:, -1]
embedding = F.normalize((hidden @ head.T).float(), dim=-1)