From e554f697bbf0090c5f1038c542cd3825017bc628 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 30 Jun 2026 12:55:10 +0200 Subject: [PATCH 01/17] [RNE Rewrite] Add text and image embeddings pipelines Add int64/Long tensor dtype support and text/image embeddings tasks, hooks, and model registry entries, plus an interactive text-embeddings demo screen in apps/nlp. Closes #1247 --- apps/nlp/app/_layout.tsx | 7 + apps/nlp/app/index.tsx | 3 + apps/nlp/app/text-embeddings/index.tsx | 410 ++++++++++++++++++ apps/nlp/package.json | 3 +- .../cpp/core/dtype.cpp | 13 +- .../react-native-executorch/cpp/core/dtype.h | 1 + .../cpp/extensions/cv/utils.h | 2 + .../src/core/tensor.ts | 10 +- .../extensions/cv/tasks/imageEmbeddings.ts | 87 ++++ .../src/extensions/nlp/index.ts | 1 + .../extensions/nlp/tasks/textEmbeddings.ts | 110 +++++ .../src/hooks/useImageEmbeddings.ts | 45 ++ .../src/hooks/useTextEmbeddings.ts | 49 +++ packages/react-native-executorch/src/index.ts | 4 + .../react-native-executorch/src/models.ts | 89 ++++ 15 files changed, 827 insertions(+), 7 deletions(-) create mode 100644 apps/nlp/app/text-embeddings/index.tsx create mode 100644 packages/react-native-executorch/src/extensions/cv/tasks/imageEmbeddings.ts create mode 100644 packages/react-native-executorch/src/extensions/nlp/tasks/textEmbeddings.ts create mode 100644 packages/react-native-executorch/src/hooks/useImageEmbeddings.ts create mode 100644 packages/react-native-executorch/src/hooks/useTextEmbeddings.ts diff --git a/apps/nlp/app/_layout.tsx b/apps/nlp/app/_layout.tsx index bdcfc39660..352cf99d44 100644 --- a/apps/nlp/app/_layout.tsx +++ b/apps/nlp/app/_layout.tsx @@ -27,6 +27,13 @@ export default function Layout() { title: 'Tokenizer', }} /> + ); } diff --git a/apps/nlp/app/index.tsx b/apps/nlp/app/index.tsx index 98ff59ab97..57b5aea6ae 100644 --- a/apps/nlp/app/index.tsx +++ b/apps/nlp/app/index.tsx @@ -14,6 +14,9 @@ export default function Home() { router.navigate('tokenizer/')}> Tokenizer + router.navigate('text-embeddings/')}> + Text Embeddings + ); diff --git a/apps/nlp/app/text-embeddings/index.tsx b/apps/nlp/app/text-embeddings/index.tsx new file mode 100644 index 0000000000..bf2f8c1b5e --- /dev/null +++ b/apps/nlp/app/text-embeddings/index.tsx @@ -0,0 +1,410 @@ +import React, { useEffect, useState } from 'react'; +import { + View, + Text, + TextInput, + ScrollView, + StyleSheet, + TouchableOpacity, + KeyboardAvoidingView, + Platform, +} from 'react-native'; +import { useTextEmbeddings, models, type TextEmbeddingsModel } from 'react-native-executorch'; +import ScreenWrapper from '../../components/ScreenWrapper'; +import { ModelStatus } from '../../components/ModelStatus'; +import { Button } from '../../components/Button'; +import { theme } from '../../theme'; + +const MODELS: { label: string; value: TextEmbeddingsModel }[] = [ + { label: 'MiniLM L6', value: models.textEmbeddings.ALL_MINILM_L6_V2 }, + { label: 'MPNet Base', value: models.textEmbeddings.ALL_MPNET_BASE_V2 }, + { label: 'MultiQA MiniLM', value: models.textEmbeddings.MULTI_QA_MINILM_L6_COS_V1 }, + { label: 'MultiQA MPNet', value: models.textEmbeddings.MULTI_QA_MPNET_BASE_DOT_V1 }, + { label: 'Paraphrase ML', value: models.textEmbeddings.PARAPHRASE_MULTILINGUAL_MINILM_L12_V2 }, + { label: 'DistilUSE ML', value: models.textEmbeddings.DISTILUSE_BASE_MULTILINGUAL_CASED_V2 }, + { label: 'CLIP Text', value: models.textEmbeddings.CLIP_VIT_BASE_PATCH32_TEXT }, +]; + +const STARTER_SENTENCES = [ + 'The weather is lovely today.', + "It's so sunny outside!", + 'He drove to the stadium.', + 'A man is eating a piece of bread.', + 'The cat sleeps on the warm windowsill.', +]; + +// True cosine similarity. Most of these models output L2-normalized embeddings +// (so a plain dot product would do), but some — e.g. DistilUSE — do not, so we +// divide by the norms to stay correct regardless of whether the model normalizes. +const cosine = (a: Float32Array, b: Float32Array) => { + let dot = 0; + let na = 0; + let nb = 0; + for (let i = 0; i < a.length; i++) { + dot += a[i]! * b[i]!; + na += a[i]! * a[i]!; + nb += b[i]! * b[i]!; + } + const denom = Math.sqrt(na) * Math.sqrt(nb); + return denom > 0 ? dot / denom : 0; +}; + +type Entry = { sentence: string; embedding: Float32Array }; +type Match = { sentence: string; similarity: number }; + +// A model switch disposes the previous task while a forward() may still be +// running; the resulting "disposed" rejection is expected and not worth showing. +const isDisposedError = (msg: string) => /disposed/i.test(msg); + +function TextEmbeddingsContent() { + const [selected, setSelected] = useState(0); + const { isReady, downloadProgress, error, forward } = useTextEmbeddings(MODELS[selected]!.value); + + const [library, setLibrary] = useState([]); + const [input, setInput] = useState(''); + const [matches, setMatches] = useState(null); + const [queryText, setQueryText] = useState(''); + const [busy, setBusy] = useState(false); + const [runError, setRunError] = useState(null); + const [embedMs, setEmbedMs] = useState(null); + + const ready = isReady && !!forward; + + // Seed the library with starter sentences whenever a model finishes loading. + // `forward`'s identity changes only when the underlying model instance does, + // so this re-seeds on every model switch. The `cancelled` flag stops a pending + // forward() from a superseded model — whose tokenizer/model the hook has since + // disposed — from writing stale state or surfacing a "disposed" error. + useEffect(() => { + if (!ready || !forward) return; + let cancelled = false; + (async () => { + setBusy(true); + setRunError(null); + try { + const entries: Entry[] = []; + for (const sentence of STARTER_SENTENCES) { + const embedding = await forward(sentence); + if (cancelled) return; + entries.push({ sentence, embedding }); + } + setLibrary(entries); + } catch (e: any) { + if (!cancelled) setRunError(e?.message ?? String(e)); + } finally { + if (!cancelled) setBusy(false); + } + })(); + return () => { + cancelled = true; + }; + }, [forward, ready]); + + const selectModel = (i: number) => { + if (i === selected) return; + setSelected(i); + setLibrary([]); + setMatches(null); + setQueryText(''); + setEmbedMs(null); + }; + + const findSimilar = async () => { + if (!forward || !input.trim() || library.length === 0) return; + setBusy(true); + setRunError(null); + try { + const start = Date.now(); + const q = await forward(input.trim()); + setEmbedMs(Date.now() - start); + const ranked = library + .map(({ sentence, embedding }) => ({ sentence, similarity: cosine(q, embedding) })) + .sort((a, b) => b.similarity - a.similarity); + setQueryText(input.trim()); + setMatches(ranked); + } catch (e: any) { + const msg = e?.message ?? String(e); + if (!isDisposedError(msg)) setRunError(msg); + } finally { + setBusy(false); + } + }; + + const addToLibrary = async () => { + if (!forward || !input.trim()) return; + setBusy(true); + setRunError(null); + try { + const start = Date.now(); + const embedding = await forward(input.trim()); + setEmbedMs(Date.now() - start); + setLibrary((prev) => [...prev, { sentence: input.trim(), embedding }]); + setInput(''); + setMatches(null); + } catch (e: any) { + const msg = e?.message ?? String(e); + if (!isDisposedError(msg)) setRunError(msg); + } finally { + setBusy(false); + } + }; + + const removeAt = (i: number) => { + setLibrary((prev) => prev.filter((_, idx) => idx !== i)); + setMatches(null); + }; + + const clearLibrary = () => { + setLibrary([]); + setMatches(null); + setQueryText(''); + }; + + return ( + + + + Text Embeddings + + Semantic search playground. Build a library of sentences, then find the ones closest in + meaning to your query using cosine similarity over the embeddings. + + + Model + + {MODELS.map((m, i) => ( + selectModel(i)} + > + + {m.label} + + + ))} + + + + + + {runError && ( + + {runError} + + )} + + + + Sentence library ({library.length}) + {library.length > 0 && ( + + Clear + + )} + + {library.length === 0 ? ( + + {ready ? 'Library is empty — add a sentence below.' : 'Waiting for the model…'} + + ) : ( + library.map((item, i) => ( + + {item.sentence} + removeAt(i)} hitSlop={8}> + ✕ + + + )) + )} + + + + Try your sentence + + +