Skip to content

Commit

Permalink
✨ Add inferenceProvider filter when listing models (#1198)
Browse files Browse the repository at this point in the history
This PR adds the missing query param `inference_provider` when listing
models (`listModels`).
  • Loading branch information
frascuchon authored Feb 14, 2025
1 parent 3e78986 commit 4f04904
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 0 deletions.
35 changes: 35 additions & 0 deletions packages/hub/src/lib/list-models.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -80,4 +80,39 @@ describe("listModels", () => {

expect(count).to.equal(10);
});

it("should search model by inference provider", async () => {
let count = 0;
for await (const entry of listModels({
search: { inferenceProviders: ["together"] },
additionalFields: ["inferenceProviderMapping"],
limit: 10,
})) {
count++;
if (Array.isArray(entry.inferenceProviderMapping)) {
expect(entry.inferenceProviderMapping.map(({ provider }) => provider)).to.include("together");
}
}

expect(count).to.equal(10);
});

it("should search model by several inference providers", async () => {
let count = 0;
const inferenceProviders = ["together", "replicate"];
for await (const entry of listModels({
search: { inferenceProviders },
additionalFields: ["inferenceProviderMapping"],
limit: 10,
})) {
count++;
if (Array.isArray(entry.inferenceProviderMapping)) {
expect(
entry.inferenceProviderMapping.filter(({ provider }) => inferenceProviders.includes(provider)).length
).toBeGreaterThan(0);
}
}

expect(count).to.equal(10);
});
});
7 changes: 7 additions & 0 deletions packages/hub/src/lib/list-models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,10 @@ export async function* listModels<
owner?: string;
task?: PipelineType;
tags?: string[];
/**
* Will search for models that have one of the inference providers in the list.
*/
inferenceProviders?: string[];
};
hubUrl?: string;
additionalFields?: T[];
Expand All @@ -84,6 +88,9 @@ export async function* listModels<
...(params?.search?.owner ? { author: params.search.owner } : undefined),
...(params?.search?.task ? { pipeline_tag: params.search.task } : undefined),
...(params?.search?.query ? { search: params.search.query } : undefined),
...(params?.search?.inferenceProviders
? { inference_provider: params.search.inferenceProviders.join(",") }
: undefined),
}),
...(params?.search?.tags?.map((tag) => ["filter", tag]) ?? []),
...MODEL_EXPAND_KEYS.map((val) => ["expand", val] satisfies [string, string]),
Expand Down

0 comments on commit 4f04904

Please sign in to comment.