Skip to content
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

feat: add chrome-ai polyfill when browser not support #12

Merged
merged 6 commits into from
Jul 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/strange-adults-give.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"chrome-ai": minor
---

feat: add chrome-ai polyfill when chrome version not ok
12 changes: 10 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@
> ⚠️ Note:
> * This module is under development and may contain errors and frequent incompatible changes.
> * Chrome's implementation of [built-in AI with Gemini Nano](https://developer.chrome.com/docs/ai/built-in) is an experiment and will change as they test and address feedback.
> * If you've never heard of it before, [follow these steps](#enable-ai-in-chrome) to turn on Chrome's built-in AI.
> * If you've never heard of it before, [follow these steps](#enabling-ai-in-chrome) to turn on Chrome's built-in AI.

## 📦 Installation

Expand Down Expand Up @@ -189,7 +189,7 @@ for await (const partialObject of result.partialObjectStream) {

> Due to model reasons, `toolCall/functionCall` are not supported. We are making an effort to implement these functions by prompt engineering.

## Enable AI in Chrome
## Enabling AI in Chrome

Chrome built-in AI is a preview feature, you need to use chrome version 127 or greater, now in [dev](https://www.google.com/chrome/dev/?extra=devchannel) or [canary](https://www.google.com/chrome/canary/) channel, [may release on stable chanel at Jul 17, 2024](https://chromestatus.com/roadmap).

Expand All @@ -198,6 +198,14 @@ After then, you should turn on these flags:
* [chrome://flags/#optimization-guide-on-device-model](chrome://flags/#optimization-guide-on-device-model): `Enabled BypassPrefRequirement`
* [chrome://components/](chrome://components/): Click `Optimization Guide On Device Model` to download the model.

Or you can try using the experimental feature: `chrome-ai/polyfill`, to use `chrome-ai` in any browser that supports WebGPU and WebAssembly.

```ts
import 'chrome-ai/polyfill';
// or
requre('chrome-ai/polyfill');
```

## License

[MIT](LICENSE) License © 2024 [Jeason](https://github.com/jeasonstudio)
21 changes: 9 additions & 12 deletions app/components/outputs.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ export const Outputs = React.forwardRef<
>(({ children, ...props }, ref) => {
let infoElement = null;

const [isBrowserSupport, setIsBrowserSupport] = React.useState(true);
const [isEnabledFlags, setIsEnabledFlags] = React.useState(true);

React.useEffect(() => {
Expand All @@ -23,7 +22,7 @@ export const Outputs = React.forwardRef<
return raw ? parseInt(raw[2], 10) : 0;
}
const version = getChromeVersion();
setIsBrowserSupport(version >= 127);
// setIsBrowserSupport(version >= 127);

setIsEnabledFlags(!!('ai' in globalThis));

Expand All @@ -32,7 +31,7 @@ export const Outputs = React.forwardRef<
});
}, []);

if (!isBrowserSupport || !isEnabledFlags) {
if (!isEnabledFlags) {
infoElement = (
<div className="w-[500px] m-auto flex flex-col gap-4">
<h2 className="scroll-m-20 border-b pb-2 text-3xl font-semibold tracking-tight first:mt-0">
Expand All @@ -43,15 +42,13 @@ export const Outputs = React.forwardRef<
experimental and will change as they test and address feedback.
</p>

{isBrowserSupport ? null : (
<Alert variant="destructive">
<AlertCircle className="h-4 w-4" />
<AlertTitle>Your browser is not supported.</AlertTitle>
<AlertDescription>
Please update Chrome to version 127 or higher.
</AlertDescription>
</Alert>
)}
<Alert variant="destructive">
<AlertCircle className="h-4 w-4" />
<AlertTitle>Your browser is not supported.</AlertTitle>
<AlertDescription>
Please update Chrome to version 127 or higher.
</AlertDescription>
</Alert>

{isEnabledFlags ? null : (
<div className="flex flex-col items-start justify-start">
Expand Down
1 change: 1 addition & 0 deletions app/orders/page.tsx
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
'use client';

import 'chrome-ai/polyfill';
import { LoaderCircle, SquareKanban } from 'lucide-react';
import { Badge } from '../components/ui/badge';
import { Button } from '../components/ui/button';
Expand Down
1 change: 1 addition & 0 deletions app/page.tsx
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
'use client';

require('chrome-ai/polyfill');
import {
CornerDownLeft,
Command,
Expand Down
17 changes: 14 additions & 3 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,23 @@
"exports": {
".": {
"import": "./dist/index.mjs",
"require": "./dist/index.js"
"require": "./dist/index.js",
"types": "./dist/index.d.ts"
},
"./polyfill": {
"import": "./dist/polyfill.mjs",
"require": "./dist/polyfill.js",
"types": "./dist/polyfill.d.ts"
}
},
"sideEffects": false,
"sideEffects": [
"./dist/polyfill.global.js",
"./dist/polyfill.js",
"./dist/polyfill.mjs"
],
"scripts": {
"dev": "tsup --tsconfig=tsconfig.tsup.json --watch",
"build": "tsup --tsconfig=tsconfig.tsup.json",
"build": "tsup --tsconfig=tsconfig.tsup.json --clean",
"test": "vitest",
"test:coverage": "vitest --coverage",
"test:ci": "vitest --run",
Expand All @@ -36,6 +46,7 @@
},
"dependencies": {
"@ai-sdk/provider": "^0.0.11",
"@mediapipe/tasks-genai": "^0.10.14",
"@mediapipe/tasks-text": "^0.10.14",
"debug": "^4.3.5"
},
Expand Down
7 changes: 7 additions & 0 deletions pnpm-lock.yaml

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

16 changes: 13 additions & 3 deletions src/global.d.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
export type ChromeAISessionAvailable = 'no' | 'readily';
export type ChromeAISessionAvailable = 'no' | 'after-download' | 'readily';

export interface ChromeAISessionOptions {
temperature?: number;
Expand All @@ -18,11 +18,21 @@ export interface ChromePromptAPI {
canCreateTextSession: () => Promise<ChromeAISessionAvailable>;
defaultGenericSessionOptions: () => Promise<ChromeAISessionOptions>;
defaultTextSessionOptions: () => Promise<ChromeAISessionOptions>;
createGenericSession: (options?: ChromeAISessionOptions) => Promise<ChromeAISession>;
createTextSession: (options?: ChromeAISessionOptions) => Promise<ChromeAISession>;
createGenericSession: (
options?: ChromeAISessionOptions
) => Promise<ChromeAISession>;
createTextSession: (
options?: ChromeAISessionOptions
) => Promise<ChromeAISession>;
}

export interface PolyfillChromeAIOptions {
llmModelAssetPath: string;
filesetBasePath: string;
}

declare global {
var ai: ChromePromptAPI;
var model = ai;
var __polyfill_ai_options__: Partial<PolyfillChromeAIOptions> | undefined;
}
1 change: 1 addition & 0 deletions src/index.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
export * from './language-model';
export * from './embedding-model';
export * from './chromeai';
export * from './polyfill/session';
3 changes: 3 additions & 0 deletions src/polyfill/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
import { polyfillChromeAI } from './session';

polyfillChromeAI(globalThis.__polyfill_ai_options__);
124 changes: 124 additions & 0 deletions src/polyfill/session.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
import {
FilesetResolver,
LlmInference,
ProgressListener,
} from '@mediapipe/tasks-genai';
import {
ChromeAISession,
ChromeAISessionAvailable,
ChromeAISessionOptions,
ChromePromptAPI,
PolyfillChromeAIOptions,
} from '../global';
import createDebug from 'debug';

const debug = createDebug('chromeai:polyfill');

class PolyfillChromeAISession implements ChromeAISession {
public constructor(private llm: LlmInference) {
debug('PolyfillChromeAISession created', llm);
}

private generateText = async (prompt: string): Promise<string> => {
const response = await this.llm.generateResponse(prompt);
debug('generateText', prompt, response);
return response;
};

private streamText = (prompt: string): ReadableStream<string> => {
debug('streamText', prompt);
const stream = new ReadableStream<string>({
start: (controller) => {
const listener: ProgressListener = (
partialResult: string,
done: boolean
) => {
controller.enqueue(partialResult);
if (done) {
controller.close();
}
};
this.llm.generateResponse(prompt, listener);
},
cancel: (reason) => {
console.warn('stream text canceled', reason);
},
});
debug('streamText', prompt);
return stream;
};

public destroy = async () => this.llm.close();
public prompt = this.generateText;
public execute = this.generateText;
public promptStreaming = this.streamText;
public executeStreaming = this.streamText;
}

/**
* Model: https://huggingface.co/oongaboongahacker/Gemini-Nano
*/
export class PolyfillChromeAI implements ChromePromptAPI {
private aiOptions: PolyfillChromeAIOptions = {
// About 1.78GB, should cache by browser
llmModelAssetPath:
'https://huggingface.co/oongaboongahacker/Gemini-Nano/resolve/main/weights.bin',
filesetBasePath: 'https://unpkg.com/@mediapipe/tasks-genai/wasm/',
};

public constructor(aiOptions: Partial<PolyfillChromeAIOptions> = {}) {
this.aiOptions = Object.assign(this.aiOptions, aiOptions);
debug('PolyfillChromeAI created', this.aiOptions);
this.modelAssetBuffer = fetch(this.aiOptions.llmModelAssetPath).then(
(response) => response.body!.getReader()
)!;
}

private modelAssetBuffer: Promise<ReadableStreamDefaultReader>;

private canCreateSession = async (): Promise<ChromeAISessionAvailable> => {
// TODO@jeasonstudio:
// * if browser do not support WebAssembly/WebGPU, return 'no';
// * check if modelAssetBuffer is downloaded, if not, return 'after-download';
return 'readily';
};
private defaultSessionOptions =
async (): Promise<ChromeAISessionOptions> => ({
temperature: 0.800000011920929,
topK: 3,
});

private createSession = async (
options?: ChromeAISessionOptions
): Promise<ChromeAISession> => {
const argv = options ?? (await this.defaultSessionOptions());
const fileset = await FilesetResolver.forGenAiTasks(
this.aiOptions.filesetBasePath
);
const llm = await LlmInference.createFromOptions(fileset, {
baseOptions: {
modelAssetBuffer: await this.modelAssetBuffer,
},
temperature: argv.temperature,
topK: argv.topK,
});
const session = new PolyfillChromeAISession(llm);
debug('createSession', options, session);
return session;
};

public canCreateGenericSession = this.canCreateSession;
public canCreateTextSession = this.canCreateSession;
public defaultGenericSessionOptions = this.defaultSessionOptions;
public defaultTextSessionOptions = this.defaultSessionOptions;
public createGenericSession = this.createSession;
public createTextSession = this.createSession;
}

export const polyfillChromeAI = (
options?: Partial<PolyfillChromeAIOptions>
) => {
const ai = new PolyfillChromeAI(options);
globalThis.ai = globalThis.ai || ai;
globalThis.model = globalThis.model || ai;
};
2 changes: 2 additions & 0 deletions tsconfig.json
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@
"extends": "./tsconfig.base.json",
"compilerOptions": {
"plugins": [{ "name": "next" }],
"baseUrl": "./",
"paths": {
"@/*": ["./*"],
"chrome-ai/*": ["./src/*"],
"chrome-ai": ["./src/index.ts"]
},
"incremental": true
Expand Down
8 changes: 8 additions & 0 deletions tsup.config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,4 +7,12 @@ export default defineConfig([
format: ['cjs', 'esm'],
sourcemap: true,
},
{
dts: true,
entry: {
polyfill: 'src/polyfill/index.ts',
},
format: ['cjs', 'esm', 'iife'],
sourcemap: true,
},
]);
1 change: 1 addition & 0 deletions vitest.config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ export default defineConfig({
// If you want a coverage reports even if your tests are failing, include the reportOnFailure option
reportOnFailure: true,
include: ['src/**/*.ts'],
exclude: ['src/polyfill/*'], // TODO@jeasonstudio: finish the polyfill tests
},
},
});