|
| 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