Skip to content

VinF Hybrid Inference: consolidate onDeviceParams initialization #8969

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

Merged
merged 4 commits into from
Apr 24, 2025
Merged
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
64 changes: 32 additions & 32 deletions packages/vertexai/src/methods/chrome-adapter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,10 @@ export class ChromeAdapter {
constructor(
private languageModelProvider?: LanguageModel,
private mode?: InferenceMode,
private onDeviceParams?: LanguageModelCreateOptions
) {}
private onDeviceParams: LanguageModelCreateOptions = {}
) {
this.addImageTypeAsExpectedInput();
}

/**
* Checks if a given request can be made on-device.
Expand All @@ -64,12 +66,8 @@ export class ChromeAdapter {
return false;
}

const availability = await this.languageModelProvider?.availability();

// Triggers async model download so it'll be available next time.
if (availability === Availability.downloadable) {
this.download();
}
// Triggers out-of-band download so model will eventually become available.
const availability = await this.downloadIfAvailable();

if (this.mode === 'only_on_device') {
return true;
Expand All @@ -91,10 +89,7 @@ export class ChromeAdapter {
* @returns {@link Response}, so we can reuse common response formatting.
*/
async generateContent(request: GenerateContentRequest): Promise<Response> {
const session = await this.createSession(
// TODO: normalize on-device params during construction.
this.onDeviceParams || {}
);
const session = await this.createSession();
// TODO: support multiple content objects when Chrome supports
// sequence<LanguageModelMessage>
const contents = await Promise.all(
Expand All @@ -115,10 +110,7 @@ export class ChromeAdapter {
async generateContentStream(
request: GenerateContentRequest
): Promise<Response> {
const session = await this.createSession(
// TODO: normalize on-device params during construction.
this.onDeviceParams || {}
);
const session = await this.createSession();
// TODO: support multiple content objects when Chrome supports
// sequence<LanguageModelMessage>
const contents = await Promise.all(
Expand Down Expand Up @@ -155,7 +147,22 @@ export class ChromeAdapter {
}

/**
* Triggers the download of an on-device model.
* Encapsulates logic to get availability and download a model if one is downloadable.
*/
private async downloadIfAvailable(): Promise<Availability | undefined> {
const availability = await this.languageModelProvider?.availability(
this.onDeviceParams
);

if (availability === Availability.downloadable) {
this.download();
}

return availability;
}

/**
* Triggers out-of-band download of an on-device model.
*
* <p>Chrome only downloads models as needed. Chrome knows a model is needed when code calls
* LanguageModel.create.</p>
Expand All @@ -168,10 +175,8 @@ export class ChromeAdapter {
return;
}
this.isDownloading = true;
const options = this.onDeviceParams || {};
ChromeAdapter.addImageTypeAsExpectedInput(options);
this.downloadPromise = this.languageModelProvider
?.create(options)
?.create(this.onDeviceParams)
.then(() => {
this.isDownloading = false;
});
Expand Down Expand Up @@ -214,19 +219,16 @@ export class ChromeAdapter {
* <p>Chrome will remove a model from memory if it's no longer in use, so this method ensures a
* new session is created before an old session is destroyed.</p>
*/
private async createSession(
// TODO: define a default value, since these are optional.
options: LanguageModelCreateOptions
): Promise<LanguageModel> {
private async createSession(): Promise<LanguageModel> {
if (!this.languageModelProvider) {
throw new AIError(
AIErrorCode.REQUEST_ERROR,
'Chrome AI requested for unsupported browser version.'
);
}
// TODO: could we use this.onDeviceParams instead of passing in options?
ChromeAdapter.addImageTypeAsExpectedInput(options);
const newSession = await this.languageModelProvider!.create(options);
const newSession = await this.languageModelProvider.create(
this.onDeviceParams
);
if (this.oldSession) {
this.oldSession.destroy();
}
Expand All @@ -235,11 +237,9 @@ export class ChromeAdapter {
return newSession;
}

private static addImageTypeAsExpectedInput(
options: LanguageModelCreateOptions
): void {
options.expectedInputs = options.expectedInputs || [];
options.expectedInputs.push({ type: 'image' });
private addImageTypeAsExpectedInput(): void {
// Defaults to support image inputs for convenience.
this.onDeviceParams.expectedInputs ??= [{ type: 'image' }];
}

/**
Expand Down
Loading