-
Notifications
You must be signed in to change notification settings - Fork 85
[RNE Rewrite] Add text and image embeddings pipelines #1292
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
msluszniak
wants to merge
11
commits into
rne-rewrite
Choose a base branch
from
@ms/add-embeddings
base: rne-rewrite
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+1,271
−41
Open
Changes from all commits
Commits
Show all changes
11 commits
Select commit
Hold shift + click to select a range
df762d6
[RNE Rewrite] Add text and image embeddings pipelines
msluszniak b9f35d8
feat: variable-length text embeddings via get_dynamic_dims
msluszniak df8d82f
Address review: trim comments, document get_dynamic_dims, add image e…
msluszniak 0f2bb51
Redesign image embeddings demo as CLIP zero-shot classifier
msluszniak ace0d99
Address review: cache dynamic dims, harden validation, fix edge cases
msluszniak 0fa2742
fix(computer-vision): add JSDoc @param/@returns for skImageToBuffer
msluszniak 6c3ccc4
feat: port get_dynamic_dims range validation onto new model validation
msluszniak da219cd
refactor(embeddings): pre-allocate static tensors via `as const` array
msluszniak 7e402b9
refactor(model): generic per-method dynamic input shapes
msluszniak d517725
refactor(embeddings): rename to Embedder/embed, drop domain task exports
msluszniak de82f79
chore(embeddings): pin all text-embedding models to v0.10.0
msluszniak File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,250 @@ | ||
| import React, { useState } from 'react'; | ||
| import { View, Text, StyleSheet, ScrollView, TextInput, TouchableOpacity } from 'react-native'; | ||
| import { useSafeAreaInsets } from 'react-native-safe-area-context'; | ||
| import { commonStyles, ColorPalette } from '../../theme'; | ||
| import { useImage } from '@shopify/react-native-skia'; | ||
| import { useImageEmbedder, useTextEmbedder, models } from 'react-native-executorch'; | ||
| import ScreenWrapper from '../../components/ScreenWrapper'; | ||
| import { getImage, skImageToBuffer } from '../../utils'; | ||
| import { ModelPicker, type ModelOption } from '../../components/ModelPicker'; | ||
| import { ImageViewport } from '../../components/ImageViewport'; | ||
| import { ModelStatus } from '../../components/ModelStatus'; | ||
| import { LatencyIndicator } from '../../components/LatencyIndicator'; | ||
| import { Button } from '../../components/Button'; | ||
|
|
||
| const IMAGE_MODEL_OPTIONS: ModelOption[] = [ | ||
| { | ||
| label: 'CLIP ViT-B/32 (INT8)', | ||
| value: models.imageEmbeddings.CLIP_VIT_BASE_PATCH32.XNNPACK_INT8, | ||
| }, | ||
| { | ||
| label: 'CLIP ViT-B/32 (FP32)', | ||
| value: models.imageEmbeddings.CLIP_VIT_BASE_PATCH32.XNNPACK_FP32, | ||
| }, | ||
| ]; | ||
|
|
||
| const DEFAULT_LABELS = [ | ||
| 'a photo of a dog', | ||
| 'a photo of a cat', | ||
| 'a landscape photo', | ||
| 'a photo of food', | ||
| 'a photo of people', | ||
| ]; | ||
|
|
||
| // CLIP text and image embeddings are L2-normalized, so their cosine similarity | ||
| // is the dot product. | ||
| const dot = (a: Float32Array, b: Float32Array) => { | ||
| let s = 0; | ||
| for (let i = 0; i < a.length; i++) { | ||
| s += a[i]! * b[i]!; | ||
| } | ||
| return s; | ||
| }; | ||
|
|
||
| function ImageEmbeddingsContent() { | ||
| const [selectedImageModel, setSelectedImageModel] = useState<any>(IMAGE_MODEL_OPTIONS[0].value); | ||
| const [imageUri, setImageUri] = useState<string | null>(null); | ||
| const [labels, setLabels] = useState<string[]>(DEFAULT_LABELS); | ||
| const [newLabel, setNewLabel] = useState(''); | ||
| const [results, setResults] = useState<{ label: string; score: number }[]>([]); | ||
| const [latency, setLatency] = useState<number | null>(null); | ||
| const [isProcessing, setIsProcessing] = useState(false); | ||
| const [error, setError] = useState<string | null>(null); | ||
|
|
||
| const insets = useSafeAreaInsets(); | ||
| const skiaImage = useImage(imageUri, (err) => setError(err.message || String(err))); | ||
|
|
||
| // Zero-shot classification pairs a CLIP image encoder with the CLIP text | ||
| // encoder and scores the image against each text label by embedding similarity. | ||
| const imageModel = useImageEmbedder(selectedImageModel); | ||
| const textModel = useTextEmbedder(models.textEmbeddings.CLIP_VIT_BASE_PATCH32_TEXT); | ||
|
|
||
| const ready = imageModel.isReady && textModel.isReady; | ||
|
|
||
| const pickImage = async () => { | ||
| setError(null); | ||
| try { | ||
| const uri = await getImage(false); | ||
| if (!uri) return; | ||
| setImageUri(uri); | ||
| setResults([]); | ||
| setLatency(null); | ||
| } catch (e: any) { | ||
| setError(e.message || String(e)); | ||
| } | ||
| }; | ||
|
|
||
| const classify = async () => { | ||
| if (!skiaImage || !ready || !imageModel.embed || !textModel.embed) return; | ||
| setIsProcessing(true); | ||
| setError(null); | ||
| try { | ||
| const start = Date.now(); | ||
| const imageEmbedding = await imageModel.embed(skImageToBuffer(skiaImage)); | ||
| const scored: { label: string; score: number }[] = []; | ||
| for (const label of labels) { | ||
| const textEmbedding = await textModel.embed(label); | ||
| scored.push({ label, score: dot(imageEmbedding, textEmbedding) }); | ||
| } | ||
| scored.sort((a, b) => b.score - a.score); | ||
| setLatency(Date.now() - start); | ||
| setResults(scored); | ||
| } catch (e: any) { | ||
| setError(e.message || String(e)); | ||
| } finally { | ||
| setIsProcessing(false); | ||
| } | ||
| }; | ||
|
|
||
| const addLabel = () => { | ||
| const trimmed = newLabel.trim(); | ||
| if (!trimmed || labels.includes(trimmed)) return; | ||
| setLabels((prev) => [...prev, trimmed]); | ||
| setNewLabel(''); | ||
| setResults([]); | ||
| }; | ||
|
|
||
| const removeLabel = (label: string) => { | ||
| setLabels((prev) => prev.filter((l) => l !== label)); | ||
| setResults((prev) => prev.filter((r) => r.label !== label)); | ||
| }; | ||
|
|
||
| const activeError = imageModel.error | ||
| ? String(imageModel.error) | ||
| : textModel.error | ||
| ? String(textModel.error) | ||
| : error; | ||
|
|
||
| return ( | ||
| <ScrollView | ||
| style={commonStyles.container} | ||
| contentContainerStyle={[commonStyles.contentContainer, { paddingBottom: insets.bottom + 24 }]} | ||
| > | ||
| <Text style={commonStyles.description}> | ||
| Pick an image, then rank text labels by how well CLIP matches them to it (zero-shot | ||
| classification). | ||
| </Text> | ||
|
|
||
| <ModelPicker | ||
| label="Image model" | ||
| options={IMAGE_MODEL_OPTIONS} | ||
| selectedValue={selectedImageModel} | ||
| onValueChange={(model) => { | ||
| setSelectedImageModel(model); | ||
| setResults([]); | ||
| setLatency(null); | ||
| }} | ||
| /> | ||
|
|
||
| <ModelStatus | ||
| isReady={ready} | ||
| downloadProgress={Math.min(imageModel.downloadProgress, textModel.downloadProgress)} | ||
| error={activeError} | ||
| modelTypeLabel="CLIP models" | ||
| /> | ||
|
|
||
| <ImageViewport skiaImage={skiaImage} onPressPlaceholder={pickImage} /> | ||
|
|
||
| <View style={commonStyles.buttonRow}> | ||
| <Button title="Pick image" onPress={pickImage} variant="secondary" /> | ||
| <Button | ||
| title="Find best label" | ||
| onPress={classify} | ||
| disabled={!skiaImage || !ready || isProcessing} | ||
| loading={isProcessing} | ||
| /> | ||
| </View> | ||
|
|
||
| <LatencyIndicator latency={latency} /> | ||
|
|
||
| {results.length > 0 && ( | ||
| <View style={styles.card}> | ||
| <Text style={styles.cardTitle}>Results</Text> | ||
| {results.map((r, i) => ( | ||
| <View key={r.label} style={styles.row}> | ||
| <Text style={[styles.rowLabel, i === 0 && styles.topLabel]} numberOfLines={1}> | ||
| {i === 0 ? '🥇 ' : ''} | ||
| {r.label} | ||
| </Text> | ||
| <Text style={styles.rowScore}>{r.score.toFixed(3)}</Text> | ||
| </View> | ||
| ))} | ||
| </View> | ||
| )} | ||
|
|
||
| <View style={styles.card}> | ||
| <Text style={styles.cardTitle}>Labels</Text> | ||
| {labels.map((label) => ( | ||
| <View key={label} style={styles.row}> | ||
| <Text style={styles.rowLabel} numberOfLines={1}> | ||
| {label} | ||
| </Text> | ||
| <TouchableOpacity onPress={() => removeLabel(label)} hitSlop={8}> | ||
| <Text style={styles.remove}>✕</Text> | ||
| </TouchableOpacity> | ||
| </View> | ||
| ))} | ||
| <View style={styles.addRow}> | ||
| <TextInput | ||
| style={styles.input} | ||
| placeholder="Add a label…" | ||
| placeholderTextColor="#94A3B8" | ||
| value={newLabel} | ||
| onChangeText={setNewLabel} | ||
| onSubmitEditing={addLabel} | ||
| returnKeyType="done" | ||
| /> | ||
| <Button title="Add" onPress={addLabel} disabled={!newLabel.trim()} variant="secondary" /> | ||
| </View> | ||
| </View> | ||
| </ScrollView> | ||
| ); | ||
| } | ||
|
|
||
| export default function ImageEmbeddingsScreen() { | ||
| return ( | ||
| <ScreenWrapper> | ||
| <ImageEmbeddingsContent /> | ||
| </ScreenWrapper> | ||
| ); | ||
| } | ||
|
|
||
| const styles = StyleSheet.create({ | ||
| card: { | ||
| width: '100%', | ||
| backgroundColor: '#fff', | ||
| borderRadius: 12, | ||
| padding: 16, | ||
| borderWidth: 1, | ||
| borderColor: '#e9ecef', | ||
| marginTop: 16, | ||
| }, | ||
| cardTitle: { | ||
| fontSize: 16, | ||
| fontWeight: '600', | ||
| color: ColorPalette.strongPrimary, | ||
| marginBottom: 8, | ||
| }, | ||
| row: { | ||
| flexDirection: 'row', | ||
| justifyContent: 'space-between', | ||
| alignItems: 'center', | ||
| paddingVertical: 8, | ||
| borderBottomWidth: 1, | ||
| borderBottomColor: '#f1f3f5', | ||
| }, | ||
| rowLabel: { fontSize: 14, color: '#334155', flex: 1, marginRight: 8 }, | ||
| topLabel: { fontWeight: '700', color: ColorPalette.strongPrimary }, | ||
| rowScore: { fontSize: 13, fontWeight: '600', color: ColorPalette.primary }, | ||
| remove: { fontSize: 16, color: '#94A3B8', paddingHorizontal: 4 }, | ||
| addRow: { flexDirection: 'row', alignItems: 'center', gap: 8, marginTop: 12 }, | ||
| input: { | ||
| flex: 1, | ||
| backgroundColor: '#f1f3f5', | ||
| borderRadius: 10, | ||
| paddingHorizontal: 12, | ||
| paddingVertical: 10, | ||
| fontSize: 14, | ||
| color: '#0F172A', | ||
| }, | ||
| }); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why was this change only added to classification when other CV screens use the same pattern?