Skip to content

Commit

Permalink
🏷️ Fix typing of request & streaming request
Browse files Browse the repository at this point in the history
  • Loading branch information
coyotte508 committed Apr 18, 2023
1 parent 1a43527 commit 24e8a27
Show file tree
Hide file tree
Showing 6 changed files with 114 additions and 25 deletions.
7 changes: 4 additions & 3 deletions packages/inference/src/HfInference.ts
Original file line number Diff line number Diff line change
@@ -1,18 +1,19 @@
import * as tasks from "./tasks";
import type { Options, RequestArgs } from "./types";
import type { DistributiveOmit } from "./utils/distributive-omit";

type Task = typeof tasks;

type TaskWithNoAccessToken = {
[key in keyof Task]: (
args: Omit<Parameters<Task[key]>[0], "accessToken">,
args: DistributiveOmit<Parameters<Task[key]>[0], "accessToken">,
options?: Parameters<Task[key]>[1]
) => ReturnType<Task[key]>;
};

type TaskWithNoAccessTokenNoModel = {
[key in keyof Task]: (
args: Omit<Parameters<Task[key]>[0], "accessToken" | "model">,
args: DistributiveOmit<Parameters<Task[key]>[0], "accessToken" | "model">,
options?: Parameters<Task[key]>[1]
) => ReturnType<Task[key]>;
};
Expand Down Expand Up @@ -51,7 +52,7 @@ export class HfInferenceEndpoint {
for (const [name, fn] of Object.entries(tasks)) {
Object.defineProperty(this, name, {
enumerable: false,
value: (params: Omit<RequestArgs, "model">, options: Options) =>
value: (params: RequestArgs, options: Options) =>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
fn({ ...params, accessToken, model: endpointUrl } as any, { ...defaultOptions, ...options }),
});
Expand Down
2 changes: 1 addition & 1 deletion packages/inference/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ export interface BaseArgs {
}

export type RequestArgs = BaseArgs &
({ data?: Blob | ArrayBuffer } | { inputs: unknown }) & {
({ data: Blob | ArrayBuffer } | { inputs: unknown }) & {
parameters?: Record<string, unknown>;
accessToken?: string;
};
15 changes: 15 additions & 0 deletions packages/inference/src/utils/distributive-omit.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
// https://dev.to/safareli/pick-omit-and-union-types-in-typescript-4nd9
// https://github.com/microsoft/TypeScript/issues/28339#issuecomment-467393437
/**
* This allows omitting keys from objects inside unions, without merging the individual components of the union.
*/

type Keys<T> = keyof T;
type DistributiveKeys<T> = T extends unknown ? Keys<T> : never;
type Omit_<T, K> = Omit<T, Extract<keyof T, K>>;

export type DistributiveOmit<T, K> = T extends unknown
? keyof Omit_<T, K> extends never
? never
: { [P in keyof Omit_<T, K>]: Omit_<T, K>[P] }
: never;
12 changes: 12 additions & 0 deletions packages/inference/test/HfInference.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -397,6 +397,18 @@ describe.concurrent(
generated_text: "a large brown and white giraffe standing in a field ",
});
});
it("request - google/flan-t5-xxl", async () => {
expect(
await hf.request({
model: "google/flan-t5-xxl",
inputs: "one plus two equals",
})
).toMatchObject([
{
generated_text: expect.any(String),
},
]);
});
it("endpoint - makes request to specified endpoint", async () => {
const ep = hf.endpoint("https://api-inference.huggingface.co/models/google/flan-t5-xxl");
const { generated_text } = await ep.textGeneration({
Expand Down
102 changes: 81 additions & 21 deletions packages/inference/test/tapes.json

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions packages/inference/test/vcr.ts
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,7 @@ async function vcr(
init: {
headers: omit(init.headers as Record<string, string>, "Authorization"),
method: init.method,
body: typeof init.body === "string" && init.body.length < 1_000 ? init.body : undefined,
},
response: {
body: isText ? new TextDecoder().decode(arrayBuffer) : "",
Expand Down

0 comments on commit 24e8a27

Please sign in to comment.