Skip to content

Commit 50e80e1

Browse files
feat: LFM2.5 text-embedding + ColBERT (MLX/XNNPACK) with prompts & MaxSim
Add the LFM2.5-Embedding-350M and LFM2.5-ColBERT-350M models, served from HuggingFace (MLX on iOS, XNNPACK on Android / iOS simulator). Text embeddings are unified into one runner and one hook: the native TextEmbeddings model returns the raw [numTokens, embeddingDim] matrix (numTokens === 1 for pooled models, the full sequence for multi-vector / late-interaction models like ColBERT), plus the input token ids. The TS layer reduces it — toVector() for the single-vector case, getTokenVectors() and maxSim() for late interaction. Models trained with asymmetric query/document prompts (LFM uses query:/ document:, ColBERT uses [Q] /[D] ) carry a "prompts" config; forward then requires a role argument ('query' | 'document') that auto-prepends the prompt. The role is type-enforced: required for prompted models, forbidden for plain ones. Also: tokenizer post_processor is now applied for text embeddings so the BOS special token is added (CLS-pooled models depend on it), and the text-to-image Encoder reads the new EmbeddingResult. Example app gains a semantic-search screen and a ColBERT late-interaction search screen demonstrating MaxSim. Authored with Claude. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent b180046 commit 50e80e1

18 files changed

Lines changed: 938 additions & 275 deletions

File tree

‎apps/text-embeddings/app/_layout.tsx‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,14 @@ export default function _layout() {
109109
headerTitleStyle: { color: ColorPalette.primary },
110110
}}
111111
/>
112+
<Drawer.Screen
113+
name="colbert/index"
114+
options={{
115+
drawerLabel: 'ColBERT search',
116+
title: 'ColBERT search',
117+
headerTitleStyle: { color: ColorPalette.primary },
118+
}}
119+
/>
112120
</Drawer>
113121
</GeneratingContext>
114122
);

‎apps/text-embeddings/app/clip-embeddings/index.tsx‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import {
1616
models,
1717
useTextEmbeddings,
1818
useImageEmbeddings,
19+
toVector,
1920
ImageEmbeddingsProps,
2021
} from 'react-native-executorch';
2122

@@ -101,7 +102,7 @@ function ClipEmbeddingsScreen() {
101102
const txtStart = Date.now();
102103
const scored: { label: string; similarity: number }[] = [];
103104
for (const label of labels) {
104-
const textEmbedding = await textModel.forward(label);
105+
const textEmbedding = toVector(await textModel.forward(label));
105106
scored.push({
106107
label,
107108
similarity: dotProduct(imageEmbedding, textEmbedding),
Lines changed: 289 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,289 @@
1+
import { useEffect, useState } from 'react';
2+
import {
3+
StyleSheet,
4+
Text,
5+
TextInput,
6+
TouchableOpacity,
7+
View,
8+
SafeAreaView,
9+
ScrollView,
10+
KeyboardAvoidingView,
11+
Platform,
12+
} from 'react-native';
13+
import { Ionicons } from '@expo/vector-icons';
14+
import { useIsFocused } from 'expo-router';
15+
import {
16+
models,
17+
useTextEmbeddings,
18+
maxSim,
19+
EmbeddingResult,
20+
} from 'react-native-executorch';
21+
import ColorPalette from '../../colors';
22+
import ErrorBanner from '../../components/ErrorBanner';
23+
24+
const colbertModel = models.text_embedding.lfm2_5_colbert_350m();
25+
26+
// The library auto-applies the model's [Q]/[D] prompts via forward(text, role).
27+
// Late-interaction MaxSim is a shipped util; the document skiplist (punctuation
28+
// token ids excluded from scoring) is the consumer's choice — these are the
29+
// LFM2.5-ColBERT skiplist ids.
30+
const SKIPLIST = [
31+
510, 511, 512, 513, 514, 515, 516, 517, 518, 519, 520, 521, 522, 523, 524,
32+
535, 536, 537, 538, 539, 540, 541, 568, 569, 570, 571, 572, 573, 600, 601,
33+
602, 603,
34+
];
35+
36+
const CORPUS: string[] = [
37+
'The forecast says heavy showers this afternoon.',
38+
"It's so sunny outside today!",
39+
'The home team scored in the final minute to win the match.',
40+
'Fans packed the stadium for the championship game.',
41+
'Simmer the tomatoes with garlic before adding the pasta.',
42+
'He whisked the eggs and folded in the melted chocolate.',
43+
'The new phone has a faster chip and a brighter screen.',
44+
'The flight to Tokyo was delayed by three hours.',
45+
'We hiked along the coast and camped near the cliffs.',
46+
];
47+
48+
const EXAMPLE_QUERIES: string[] = [
49+
"What's the weather like?",
50+
'Who won the match?',
51+
'How do I cook dinner?',
52+
'Tell me about the latest technology',
53+
];
54+
55+
type Ranked = { sentence: string; score: number };
56+
57+
export default function ColbertScreenWrapper() {
58+
return useIsFocused() ? <ColbertScreen /> : null;
59+
}
60+
61+
function ColbertScreen() {
62+
const model = useTextEmbeddings({ model: colbertModel });
63+
const [error, setError] = useState<string | null>(null);
64+
const [query, setQuery] = useState('');
65+
const [docEncs, setDocEncs] = useState<
66+
{ sentence: string; enc: EmbeddingResult }[]
67+
>([]);
68+
const [results, setResults] = useState<Ranked[]>([]);
69+
const [indexing, setIndexing] = useState(false);
70+
const [encodeTime, setEncodeTime] = useState<number | null>(null);
71+
72+
useEffect(
73+
() => {
74+
let cancelled = false;
75+
const indexCorpus = async () => {
76+
if (!model.isReady) return;
77+
setIndexing(true);
78+
setResults([]);
79+
try {
80+
const encs = [];
81+
for (const sentence of CORPUS) {
82+
const enc = await model.forward(sentence, 'document');
83+
if (cancelled) return;
84+
encs.push({ sentence, enc });
85+
}
86+
setDocEncs(encs);
87+
} catch (e) {
88+
setError(e instanceof Error ? e.message : String(e));
89+
} finally {
90+
if (!cancelled) setIndexing(false);
91+
}
92+
};
93+
indexCorpus();
94+
return () => {
95+
cancelled = true;
96+
};
97+
},
98+
// eslint-disable-next-line react-hooks/exhaustive-deps
99+
[model.isReady]
100+
);
101+
102+
const runSearch = async (queryText: string = query) => {
103+
const q = queryText.trim();
104+
if (!model.isReady || !q || docEncs.length === 0) return;
105+
setQuery(queryText);
106+
try {
107+
const start = Date.now();
108+
const qEnc = await model.forward(q, 'query');
109+
setEncodeTime(Date.now() - start);
110+
const ranked = docEncs
111+
.map(({ sentence, enc }) => ({
112+
sentence,
113+
score: maxSim(qEnc, enc, SKIPLIST),
114+
}))
115+
.sort((a, b) => b.score - a.score);
116+
setResults(ranked);
117+
} catch (e) {
118+
setError(e instanceof Error ? e.message : String(e));
119+
}
120+
};
121+
122+
const ready = model.isReady && !indexing && docEncs.length > 0;
123+
const canSearch = ready && !!query.trim();
124+
125+
const statusText = model.error
126+
? `Error: ${model.error}`
127+
: !model.isReady
128+
? `Loading model ${(model.downloadProgress * 100).toFixed(0)}%`
129+
: indexing
130+
? 'Indexing corpus…'
131+
: 'Ready';
132+
133+
return (
134+
<SafeAreaView style={styles.container}>
135+
<KeyboardAvoidingView
136+
style={styles.flex}
137+
behavior={Platform.OS === 'ios' ? 'padding' : undefined}
138+
>
139+
<ScrollView contentContainerStyle={styles.scroll}>
140+
<Text style={styles.heading}>ColBERT Late-Interaction Search</Text>
141+
<Text style={styles.status}>{statusText}</Text>
142+
<ErrorBanner message={error} onDismiss={() => setError(null)} />
143+
144+
<View style={styles.card}>
145+
<Text style={styles.sectionTitle}>
146+
Search the corpus ({CORPUS.length} sentences)
147+
</Text>
148+
<Text style={styles.hint}>
149+
Per-token vectors scored with MaxSim. Tap an example or type a
150+
query.
151+
</Text>
152+
<View style={styles.chipRow}>
153+
{EXAMPLE_QUERIES.map((q) => (
154+
<TouchableOpacity
155+
key={q}
156+
style={[styles.chip, !ready && styles.chipDisabled]}
157+
disabled={!ready}
158+
onPress={() => runSearch(q)}
159+
>
160+
<Text style={styles.chipText}>{q}</Text>
161+
</TouchableOpacity>
162+
))}
163+
</View>
164+
<TextInput
165+
placeholder="Type a search query..."
166+
placeholderTextColor="#94A3B8"
167+
style={styles.input}
168+
value={query}
169+
onChangeText={setQuery}
170+
onSubmitEditing={() => runSearch()}
171+
returnKeyType="search"
172+
/>
173+
<TouchableOpacity
174+
onPress={() => runSearch()}
175+
style={[styles.button, !canSearch && styles.buttonDisabled]}
176+
disabled={!canSearch}
177+
>
178+
<Ionicons
179+
name="search"
180+
size={16}
181+
color={!canSearch ? 'gray' : 'white'}
182+
/>
183+
<Text style={[styles.buttonText, !canSearch && styles.buttonTextDisabled]}>
184+
{indexing ? 'Indexing…' : 'Search'}
185+
</Text>
186+
</TouchableOpacity>
187+
{encodeTime !== null && (
188+
<Text style={styles.stats}>Query encoded in {encodeTime} ms</Text>
189+
)}
190+
</View>
191+
192+
{results.length > 0 && (
193+
<View style={styles.card}>
194+
<Text style={styles.sectionTitle}>Results</Text>
195+
{results.map((r, i) => (
196+
<View key={i} style={styles.resultRow}>
197+
<View style={styles.resultHeader}>
198+
<Text style={styles.resultText}>{r.sentence}</Text>
199+
<Text style={styles.resultScore}>{r.score.toFixed(2)}</Text>
200+
</View>
201+
<View style={styles.barTrack}>
202+
<View
203+
style={[
204+
styles.barFill,
205+
{
206+
width: `${Math.round(
207+
(results[0].score > 0 ? r.score / results[0].score : 0) * 100
208+
)}%`,
209+
},
210+
i === 0 && styles.barFillTop,
211+
]}
212+
/>
213+
</View>
214+
</View>
215+
))}
216+
</View>
217+
)}
218+
</ScrollView>
219+
</KeyboardAvoidingView>
220+
</SafeAreaView>
221+
);
222+
}
223+
224+
const styles = StyleSheet.create({
225+
container: { flex: 1, backgroundColor: '#F8FAFC' },
226+
flex: { flex: 1 },
227+
scroll: { padding: 20 },
228+
heading: { fontSize: 22, fontWeight: '500', marginBottom: 8, color: '#0F172A' },
229+
status: { fontSize: 14, color: '#64748B', marginBottom: 12 },
230+
card: {
231+
backgroundColor: '#fff',
232+
padding: 16,
233+
borderRadius: 16,
234+
borderColor: '#E2E8F0',
235+
borderWidth: 2,
236+
marginBottom: 20,
237+
},
238+
sectionTitle: { fontSize: 16, fontWeight: '500', marginBottom: 8, color: '#1E293B' },
239+
hint: { fontSize: 13, color: '#64748B', marginBottom: 12, lineHeight: 18 },
240+
chipRow: { flexDirection: 'row', flexWrap: 'wrap', gap: 8, marginBottom: 12 },
241+
chip: {
242+
backgroundColor: '#EEF2FF',
243+
borderColor: '#C7D2FE',
244+
borderWidth: 1,
245+
borderRadius: 16,
246+
paddingHorizontal: 12,
247+
paddingVertical: 6,
248+
},
249+
chipDisabled: { opacity: 0.4 },
250+
chipText: { fontSize: 13, color: 'navy' },
251+
input: {
252+
backgroundColor: '#F1F5F9',
253+
borderRadius: 10,
254+
padding: 10,
255+
marginBottom: 10,
256+
fontSize: 16,
257+
color: '#0F172A',
258+
minHeight: 40,
259+
},
260+
button: {
261+
backgroundColor: 'navy',
262+
borderRadius: 10,
263+
paddingVertical: 12,
264+
flexDirection: 'row',
265+
alignItems: 'center',
266+
justifyContent: 'center',
267+
},
268+
buttonDisabled: { backgroundColor: '#f0f0f0' },
269+
buttonText: { color: '#fff', fontWeight: '500', marginLeft: 6 },
270+
buttonTextDisabled: { color: 'gray' },
271+
stats: { fontSize: 13, color: '#64748B', marginTop: 8, textAlign: 'center' },
272+
resultRow: { marginBottom: 14 },
273+
resultHeader: {
274+
flexDirection: 'row',
275+
justifyContent: 'space-between',
276+
marginBottom: 6,
277+
gap: 8,
278+
},
279+
resultText: { flex: 1, fontSize: 14, color: '#334155' },
280+
resultScore: {
281+
fontSize: 14,
282+
fontWeight: '600',
283+
color: '#0F172A',
284+
fontVariant: ['tabular-nums'],
285+
},
286+
barTrack: { height: 8, borderRadius: 4, backgroundColor: '#E2E8F0', overflow: 'hidden' },
287+
barFill: { height: '100%', borderRadius: 4, backgroundColor: '#94A3B8' },
288+
barFillTop: { backgroundColor: 'navy' },
289+
});

0 commit comments

Comments
 (0)