Skip to content

Commit

Permalink
Merge pull request #17 from jeasonstudio/chore-change-model-url
Browse files Browse the repository at this point in the history
chore: update dependencies for mediapipe to version 0.10.14
  • Loading branch information
jeasonstudio authored Jul 16, 2024
2 parents 8c9d7bd + 3036a49 commit 65443a2
Show file tree
Hide file tree
Showing 8 changed files with 234 additions and 67 deletions.
5 changes: 5 additions & 0 deletions .changeset/clever-dancers-sell.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"chrome-ai": patch
---

chore: use mediapipe version 0.10.14 and change model asset cdn url
26 changes: 25 additions & 1 deletion app/components/layout.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@ import {
Github,
Twitter,
SquareKanban,
FileInput
FileInput,
FileDiff,
} from 'lucide-react';
import { Button } from '../components/ui/button';
import {
Expand Down Expand Up @@ -105,6 +106,29 @@ export const Layout: React.FC<React.PropsWithChildren<LayoutProps>> = ({
</TooltipContent>
</Tooltip>
</TooltipProvider>
<TooltipProvider delayDuration={0}>
<Tooltip>
<TooltipTrigger asChild>
<Button
variant="ghost"
size="icon"
className={cn(
'rounded-lg',
pathname === '/embedding' && 'bg-muted'
)}
aria-label="Embedding Playground"
asChild
>
<Link href="/embedding">
<FileDiff className="size-5" />
</Link>
</Button>
</TooltipTrigger>
<TooltipContent side="right" sideOffset={5}>
Embedding Playground
</TooltipContent>
</Tooltip>
</TooltipProvider>
</nav>
{/* <nav className="mt-auto grid gap-1 p-2">
<ThemeSwitcher />
Expand Down
121 changes: 121 additions & 0 deletions app/embedding/page.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
'use client';

import React from 'react';
import { Layout } from '../components/layout';
import { useForm } from 'react-hook-form';
import { z } from 'zod';
import { zodResolver } from '@hookform/resolvers/zod';
import {
Form,
FormControl,
FormField,
FormItem,
FormLabel,
FormMessage,
} from '../components/ui/form';
import { cosineSimilarity, embedMany } from 'ai';
import { Input } from '../components/ui/input';
import { Button } from '../components/ui/button';
import { Outputs } from '../components/outputs';
import { chromeai, ChromeAIEmbeddingModel } from 'chrome-ai';

const schema = z.object({
text1: z.string(),
text2: z.string(),
});

const EmbeddingPage: React.FC<unknown> = () => {
const modelRef = React.useRef<ChromeAIEmbeddingModel>();
React.useEffect(() => {
modelRef.current = chromeai('embedding');
}, []);

const form = useForm<z.infer<typeof schema>>({
resolver: zodResolver(schema),
defaultValues: {
text1: 'sunny day at the beach',
text2: 'rainy afternoon in the city',
},
});

const [cost, setCost] = React.useState<number>(0);
const [similarity, setSimilarity] = React.useState<number>(0);

async function onSubmit(data: z.infer<typeof schema>) {
setSimilarity(0);
setCost(0);
const startTime = performance.now();
const { embeddings } = await embedMany({
model: modelRef.current!,
values: [data.text1, data.text2],
});

const similarity = cosineSimilarity(embeddings[0], embeddings[1]);
setSimilarity(similarity);
setCost(performance.now() - startTime);
}
return (
<Layout>
<main className="grid flex-1 gap-4 overflow-auto p-4 md:grid-cols-2 lg:grid-cols-3 h-screen">
<Outputs className="w-full">
<h2 className="mt-4 scroll-m-20 text-3xl font-semibold tracking-tight transition-colors">
Embedding
</h2>
<p className="leading-7 mt-6">
After embedding values, you can calculate the similarity between
them using the cosineSimilarity function. This is useful to e.g.
find similar words or phrases in a dataset. You can also rank and
filter related items based on their similarity.
</p>

<Form {...form}>
<form
onSubmit={form.handleSubmit(onSubmit)}
className="w-2/3 space-y-6 mt-5"
>
<FormField
control={form.control}
name="text1"
render={({ field }) => (
<FormItem>
<FormLabel>Text 1</FormLabel>
<FormControl>
<Input placeholder="text 1" {...field} />
</FormControl>
<FormMessage />
</FormItem>
)}
/>
<FormField
control={form.control}
name="text2"
render={({ field }) => (
<FormItem>
<FormLabel>Text 2</FormLabel>
<FormControl>
<Input placeholder="text 2" {...field} />
</FormControl>
<FormMessage />
</FormItem>
)}
/>
<Button type="submit">Get Similarity</Button>
</form>
</Form>

{!!similarity && !!cost ? (
<div className="mt-6">
<div className="flex ">
<div>
Similarity: {similarity}{' '}
<span className="gray">(cost: {cost}ms)</span>
</div>
</div>
</div>
) : null}
</Outputs>
</main>
</Layout>
);
};
export default EmbeddingPage;
4 changes: 2 additions & 2 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,8 @@
},
"dependencies": {
"@ai-sdk/provider": "^0.0.11",
"@mediapipe/tasks-genai": "^0.10.14",
"@mediapipe/tasks-text": "^0.10.14",
"@mediapipe/tasks-genai": "0.10.14",
"@mediapipe/tasks-text": "0.10.14",
"debug": "^4.3.5"
},
"devDependencies": {
Expand Down
4 changes: 2 additions & 2 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

51 changes: 32 additions & 19 deletions src/embedding-model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,18 @@ import { TextEmbedder, FilesetResolver } from '@mediapipe/tasks-text';
export interface ChromeAIEmbeddingModelSettings {
/**
* An optional base path to specify the directory the Wasm files should be loaded from.
* @default 'https://pub-ddcfe353995744e89b8002f16bf98575.r2.dev/text_wasm_internal.js'
*/
wasmLoaderPath?: string;
/**
* It's about 6mb before gzip.
* @default 'https://unpkg.com/@mediapipe/tasks-text/wasm/'
* @default 'https://pub-ddcfe353995744e89b8002f16bf98575.r2.dev/text_wasm_internal.wasm'
*/
filesetBasePath?: string;
wasmBinaryPath?: string;
/**
* The model path to the model asset file.
* It's about 6.1mb before gzip.
* @default 'https://storage.googleapis.com/mediapipe-models/text_embedder/universal_sentence_encoder/float32/1/universal_sentence_encoder.tflite'
* @default 'https://pub-ddcfe353995744e89b8002f16bf98575.r2.dev/universal_sentence_encoder.tflite'
*/
modelAssetPath?: string;
/**
Expand Down Expand Up @@ -47,32 +51,41 @@ export class ChromeAIEmbeddingModel implements EmbeddingModelV1<string> {
readonly maxEmbeddingsPerCall = undefined;

private settings: ChromeAIEmbeddingModelSettings = {
filesetBasePath: 'https://unpkg.com/@mediapipe/tasks-text/wasm/',
wasmLoaderPath:
'https://pub-ddcfe353995744e89b8002f16bf98575.r2.dev/text_wasm_internal.js',
wasmBinaryPath:
'https://pub-ddcfe353995744e89b8002f16bf98575.r2.dev/text_wasm_internal.wasm',
modelAssetPath:
'https://storage.googleapis.com/mediapipe-models/text_embedder/universal_sentence_encoder/float32/1/universal_sentence_encoder.tflite',
'https://pub-ddcfe353995744e89b8002f16bf98575.r2.dev/universal_sentence_encoder.tflite',
l2Normalize: false,
quantize: false,
};
private textEmbedder: TextEmbedder | null = null;
private modelAssetBuffer!: Promise<ReadableStreamDefaultReader>;
private textEmbedder!: Promise<TextEmbedder>;

public constructor(settings: ChromeAIEmbeddingModelSettings = {}) {
this.settings = { ...this.settings, ...settings };
this.modelAssetBuffer = fetch(this.settings.modelAssetPath!).then(
(response) => response.body!.getReader()
)!;
this.textEmbedder = this.getTextEmbedder();
}

protected getTextEmbedder = async (): Promise<TextEmbedder> => {
if (this.textEmbedder !== null) return this.textEmbedder;
const textFiles = await FilesetResolver.forTextTasks(
this.settings.filesetBasePath
);
this.textEmbedder = await TextEmbedder.createFromOptions(textFiles, {
baseOptions: {
modelAssetPath: this.settings.modelAssetPath,
delegate: this.settings.delegate,
return TextEmbedder.createFromOptions(
{
wasmBinaryPath: this.settings.wasmBinaryPath!,
wasmLoaderPath: this.settings.wasmLoaderPath!,
},
l2Normalize: this.settings.l2Normalize,
quantize: this.settings.quantize,
});
return this.textEmbedder;
{
baseOptions: {
modelAssetBuffer: await this.modelAssetBuffer,
delegate: this.settings.delegate,
},
l2Normalize: this.settings.l2Normalize,
quantize: this.settings.quantize,
}
);
};

public doEmbed = async (options: {
Expand All @@ -83,7 +96,7 @@ export class ChromeAIEmbeddingModel implements EmbeddingModelV1<string> {
rawResponse?: Record<PropertyKey, any>;
}> => {
// if (options.abortSignal) console.warn('abortSignal is not supported');
const embedder = await this.getTextEmbedder();
const embedder = await this.textEmbedder;
const embeddings = options.values.map((text) => {
const embedderResult = embedder.embed(text);
const [embedding] = embedderResult.embeddings;
Expand Down
5 changes: 3 additions & 2 deletions src/global.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,9 @@ export interface ChromePromptAPI {
}

export interface PolyfillChromeAIOptions {
llmModelAssetPath: string;
filesetBasePath: string;
modelAssetPath: string;
wasmLoaderPath: string;
wasmBinaryPath: string;
}

declare global {
Expand Down
Loading

0 comments on commit 65443a2

Please sign in to comment.