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

[Inference] [Providers] Enforce task in mapping + expose them #1109

Merged
merged 8 commits into from
Jan 17, 2025
Merged
Changes from 1 commit
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
Prev Previous commit
Next Next commit
lint + format
SBrandeis committed Jan 16, 2025
commit d309548487fa721c3976e03e14b1dbcb929620ff
29 changes: 17 additions & 12 deletions packages/inference/src/lib/makeRequestOptions.ts
Original file line number Diff line number Diff line change
@@ -64,12 +64,12 @@ export async function makeRequestOptions(
? endpointUrl + `/v1/chat/completions`
: endpointUrl
: makeUrl({
model,
provider: provider ?? "hf-inference",
taskHint,
chatCompletion: chatCompletion ?? false,
forceTask,
});
model,
provider: provider ?? "hf-inference",
taskHint,
chatCompletion: chatCompletion ?? false,
forceTask,
});

const headers: Record<string, string> = {};
if (accessToken) {
@@ -122,24 +122,29 @@ export async function makeRequestOptions(
body: binary
? args.data
: JSON.stringify({
...otherArgs,
...(chatCompletion || provider === "together" ? { model } : undefined),
}),
...otherArgs,
...(chatCompletion || provider === "together" ? { model } : undefined),
}),
...(credentials ? { credentials } : undefined),
signal: options?.signal,
};

return { url, info };
}

function mapModel(params: { model: string; provider: InferenceProvider; taskHint: InferenceTask | undefined; chatCompletion: boolean | undefined }): string {
function mapModel(params: {
model: string;
provider: InferenceProvider;
taskHint: InferenceTask | undefined;
chatCompletion: boolean | undefined;
}): string {
if (params.provider === "hf-inference") {
return params.model;
}
if (!params.taskHint) {
throw new Error("taskHint must be specified when using a third-party provider")
throw new Error("taskHint must be specified when using a third-party provider");
}
const task = params.taskHint === "text-generation" && params.chatCompletion ? "conversational" : params.taskHint
const task = params.taskHint === "text-generation" && params.chatCompletion ? "conversational" : params.taskHint;
const model = (() => {
switch (params.provider) {
case "fal-ai":
3 changes: 1 addition & 2 deletions packages/inference/src/providers/fal-ai.ts
Original file line number Diff line number Diff line change
@@ -4,13 +4,12 @@ export const FAL_AI_API_BASE_URL = "https://fal.run";

type FalAiId = string;


export const FAL_AI_MODEL_IDS: ProviderMapping<FalAiId> = {
"text-to-image": {
"black-forest-labs/FLUX.1-schnell": "fal-ai/flux/schnell",
"black-forest-labs/FLUX.1-dev": "fal-ai/flux/dev",
},
"automatic-speech-recognition": {
"openai/whisper-large-v3": "fal-ai/whisper",
}
},
};
2 changes: 1 addition & 1 deletion packages/inference/src/providers/replicate.ts
Original file line number Diff line number Diff line change
@@ -9,5 +9,5 @@ export const REPLICATE_MODEL_IDS: ProviderMapping<ReplicateId> = {
"black-forest-labs/FLUX.1-schnell": "black-forest-labs/flux-schnell",
"ByteDance/SDXL-Lightning":
"bytedance/sdxl-lightning-4step:5599ed30703defd1d160a25a63321b4dec97101d98b4674bcc56e41f62f35637",
}
},
};
4 changes: 2 additions & 2 deletions packages/inference/src/providers/sambanova.ts
Original file line number Diff line number Diff line change
@@ -6,7 +6,7 @@ type SambanovaId = string;

export const SAMBANOVA_MODEL_IDS: ProviderMapping<SambanovaId> = {
/** Chat completion / conversational */
"conversational": {
conversational: {
"Qwen/Qwen2.5-Coder-32B-Instruct": "Qwen2.5-Coder-32B-Instruct",
"Qwen/Qwen2.5-72B-Instruct": "Qwen2.5-72B-Instruct",
"Qwen/QwQ-32B-Preview": "QwQ-32B-Preview",
@@ -19,5 +19,5 @@ export const SAMBANOVA_MODEL_IDS: ProviderMapping<SambanovaId> = {
"meta-llama/Llama-3.1-70B-Instruct": "Meta-Llama-3.1-70B-Instruct",
"meta-llama/Llama-3.1-405B-Instruct": "Meta-Llama-3.1-405B-Instruct",
"meta-llama/Llama-Guard-3-8B": "Meta-Llama-Guard-3-8B",
}
},
};
4 changes: 2 additions & 2 deletions packages/inference/src/providers/together.ts
Original file line number Diff line number Diff line change
@@ -19,7 +19,7 @@ export const TOGETHER_MODEL_IDS: ProviderMapping<TogetherId> = {
"black-forest-labs/FLUX.1-schnell": "black-forest-labs/FLUX.1-pro",
"stabilityai/stable-diffusion-xl-base-1.0": "stabilityai/stable-diffusion-xl-base-1.0",
},
"conversational": {
conversational: {
"databricks/dbrx-instruct": "databricks/dbrx-instruct",
"deepseek-ai/deepseek-llm-67b-chat": "deepseek-ai/deepseek-llm-67b-chat",
"google/gemma-2-9b-it": "google/gemma-2-9b-it",
@@ -54,5 +54,5 @@ export const TOGETHER_MODEL_IDS: ProviderMapping<TogetherId> = {
"text-generation": {
"meta-llama/Meta-Llama-3-8B": "meta-llama/Meta-Llama-3-8B",
"mistralai/Mixtral-8x7B-v0.1": "mistralai/Mixtral-8x7B-v0.1",
}
},
};
4 changes: 3 additions & 1 deletion packages/inference/src/providers/types.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import type { InferenceTask, ModelId } from "../types";

export type ProviderMapping<ProviderId extends string> = Partial<Record<InferenceTask | "conversational", Partial<Record<ModelId, ProviderId>>>>;
export type ProviderMapping<ProviderId extends string> = Partial<
Record<InferenceTask | "conversational", Partial<Record<ModelId, ProviderId>>>
>;
8 changes: 5 additions & 3 deletions packages/inference/test/HfInference.spec.ts
Original file line number Diff line number Diff line change
@@ -919,9 +919,11 @@ describe.concurrent("HfInference", () => {
model: "black-forest-labs/Flux.1-dev",
provider: "together",
messages: [{ role: "user", content: "Complete this sentence with words, one plus one is equal " }],
accessToken: env.HF_TOGETHER_KEY
accessToken: env.HF_TOGETHER_KEY,
})
).rejects.toThrowError("Model black-forest-labs/Flux.1-dev is not supported for task conversational and provider together")
).rejects.toThrowError(
"Model black-forest-labs/Flux.1-dev is not supported for task conversational and provider together"
);
});
})
});
});