feat(hub): adding revision property to modelInfo, datasetInfo, spaceInfo (#999)

axel7083 · web-flow · commit d623bbd6f4b4 · 2024-10-30T14:07:44.000+01:00
When getting modelInfo, datasetInfo, spaceInfo, we might need to get the
details for a specific revision.

This PR adds an optional `revision` property to those methods
diff --git a/packages/hub/src/lib/dataset-info.spec.ts b/packages/hub/src/lib/dataset-info.spec.ts
@@ -1,5 +1,7 @@
 import { describe, expect, it } from "vitest";
 import { datasetInfo } from "./dataset-info";
+import type { DatasetEntry } from "./list-datasets";
+import type { ApiDatasetInfo } from "../types/api/api-dataset";
 
 describe("datasetInfo", () => {
 	it("should return the dataset info", async () => {
@@ -16,4 +18,39 @@ describe("datasetInfo", () => {
 			private: false,
 		});
 	});
+
+	it("should return the dataset info with author", async () => {
+		const info: DatasetEntry & Pick<ApiDatasetInfo, 'author'> = await datasetInfo({
+			name: "nyu-mll/glue",
+			additionalFields: ['author'],
+		});
+		expect(info).toEqual({
+			id: "621ffdd236468d709f181e3f",
+			downloads: expect.any(Number),
+			gated: false,
+			name: "nyu-mll/glue",
+			updatedAt: expect.any(Date),
+			likes: expect.any(Number),
+			private: false,
+			author: 'nyu-mll'
+		});
+	});
+
+	it("should return the dataset info for a specific revision", async () => {
+		const info: DatasetEntry & Pick<ApiDatasetInfo, 'sha'> = await datasetInfo({
+			name: "nyu-mll/glue",
+			revision: "cb2099c76426ff97da7aa591cbd317d91fb5fcb7",
+			additionalFields: ["sha"],
+		});
+		expect(info).toEqual({
+			id: "621ffdd236468d709f181e3f",
+			downloads: expect.any(Number),
+			gated: false,
+			name: "nyu-mll/glue",
+			updatedAt: expect.any(Date),
+			likes: expect.any(Number),
+			private: false,
+			sha: 'cb2099c76426ff97da7aa591cbd317d91fb5fcb7'
+		});
+	});
 });
diff --git a/packages/hub/src/lib/dataset-info.ts b/packages/hub/src/lib/dataset-info.ts
@@ -13,6 +13,10 @@ export async function datasetInfo<
 		name: string;
 		hubUrl?: string;
 		additionalFields?: T[];
+		/**
+		 * An optional Git revision id which can be a branch name, a tag, or a commit hash.
+		 */
+		revision?: string;
 		/**
 		 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
 		 */
@@ -27,7 +31,7 @@ export async function datasetInfo<
 	]).toString();
 
 	const response = await (params.fetch || fetch)(
-		`${params?.hubUrl || HUB_URL}/api/datasets/${params.name}?${search.toString()}`,
+		`${params?.hubUrl || HUB_URL}/api/datasets/${params.name}/revision/${encodeURIComponent(params.revision ?? "HEAD")}?${search.toString()}`,
 		{
 			headers: {
 				...(accessToken ? { Authorization: `Bearer ${accessToken}` } : {}),
diff --git a/packages/hub/src/lib/model-info.spec.ts b/packages/hub/src/lib/model-info.spec.ts
@@ -1,10 +1,28 @@
 import { describe, expect, it } from "vitest";
 import { modelInfo } from "./model-info";
+import type { ModelEntry } from "./list-models";
+import type { ApiModelInfo } from "../types/api/api-model";
 
 describe("modelInfo", () => {
 	it("should return the model info", async () => {
 		const info = await modelInfo({
 			name: "openai-community/gpt2",
+		});
+		expect(info).toEqual({
+			id: "621ffdc036468d709f17434d",
+			downloads: expect.any(Number),
+			gated: false,
+			name: "openai-community/gpt2",
+			updatedAt: expect.any(Date),
+			likes: expect.any(Number),
+			task: "text-generation",
+			private: false,
+		});
+	});
+
+	it("should return the model info with author", async () => {
+		const info: ModelEntry & Pick<ApiModelInfo, 'author'> = await modelInfo({
+			name: "openai-community/gpt2",
 			additionalFields: ["author"],
 		});
 		expect(info).toEqual({
@@ -19,4 +37,23 @@ describe("modelInfo", () => {
 			private: false,
 		});
 	});
+
+	it("should return the model info for a specific revision", async () => {
+		const info: ModelEntry & Pick<ApiModelInfo, 'sha'> = await modelInfo({
+			name: "openai-community/gpt2",
+			additionalFields: ["sha"],
+			revision: 'f27b190eeac4c2302d24068eabf5e9d6044389ae',
+		});
+		expect(info).toEqual({
+			id: "621ffdc036468d709f17434d",
+			downloads: expect.any(Number),
+			gated: false,
+			name: "openai-community/gpt2",
+			updatedAt: expect.any(Date),
+			likes: expect.any(Number),
+			task: "text-generation",
+			private: false,
+			sha: "f27b190eeac4c2302d24068eabf5e9d6044389ae",
+		});
+	});
 });
diff --git a/packages/hub/src/lib/model-info.ts b/packages/hub/src/lib/model-info.ts
@@ -13,6 +13,10 @@ export async function modelInfo<
 		name: string;
 		hubUrl?: string;
 		additionalFields?: T[];
+		/**
+		 * An optional Git revision id which can be a branch name, a tag, or a commit hash.
+		 */
+		revision?: string;
 		/**
 		 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
 		 */
@@ -27,7 +31,7 @@ export async function modelInfo<
 	]).toString();
 
 	const response = await (params.fetch || fetch)(
-		`${params?.hubUrl || HUB_URL}/api/models/${params.name}?${search.toString()}`,
+		`${params?.hubUrl || HUB_URL}/api/models/${params.name}/revision/${encodeURIComponent(params.revision ?? "HEAD")}?${search.toString()}`,
 		{
 			headers: {
 				...(accessToken ? { Authorization: `Bearer ${accessToken}` } : {}),
diff --git a/packages/hub/src/lib/space-info.spec.ts b/packages/hub/src/lib/space-info.spec.ts
@@ -1,5 +1,7 @@
 import { describe, expect, it } from "vitest";
 import { spaceInfo } from "./space-info";
+import type { SpaceEntry } from "./list-spaces";
+import type { ApiSpaceInfo } from "../types/api/api-space";
 
 describe("spaceInfo", () => {
 	it("should return the space info", async () => {
@@ -15,4 +17,37 @@ describe("spaceInfo", () => {
 			sdk: "static",
 		});
 	});
+
+	it("should return the space info with author", async () => {
+		const info: SpaceEntry & Pick<ApiSpaceInfo, 'author'> = await spaceInfo({
+			name: "huggingfacejs/client-side-oauth",
+			additionalFields: ['author'],
+		});
+		expect(info).toEqual({
+			id: "659835e689010f9c7aed608d",
+			name: "huggingfacejs/client-side-oauth",
+			updatedAt: expect.any(Date),
+			likes: expect.any(Number),
+			private: false,
+			sdk: "static",
+			author: 'huggingfacejs',
+		});
+	});
+
+	it("should return the space info for a given revision", async () => {
+		const info: SpaceEntry & Pick<ApiSpaceInfo, 'sha'> = await spaceInfo({
+			name: "huggingfacejs/client-side-oauth",
+			additionalFields: ['sha'],
+			revision: 'e410a9ff348e6bed393b847711e793282d7c672e'
+		});
+		expect(info).toEqual({
+			id: "659835e689010f9c7aed608d",
+			name: "huggingfacejs/client-side-oauth",
+			updatedAt: expect.any(Date),
+			likes: expect.any(Number),
+			private: false,
+			sdk: "static",
+			sha: 'e410a9ff348e6bed393b847711e793282d7c672e',
+		});
+	});
 });
diff --git a/packages/hub/src/lib/space-info.ts b/packages/hub/src/lib/space-info.ts
@@ -14,6 +14,10 @@ export async function spaceInfo<
 		name: string;
 		hubUrl?: string;
 		additionalFields?: T[];
+		/**
+		 * An optional Git revision id which can be a branch name, a tag, or a commit hash.
+		 */
+		revision?: string;
 		/**
 		 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
 		 */
@@ -28,7 +32,7 @@ export async function spaceInfo<
 	]).toString();
 
 	const response = await (params.fetch || fetch)(
-		`${params?.hubUrl || HUB_URL}/api/spaces/${params.name}?${search.toString()}`,
+		`${params?.hubUrl || HUB_URL}/api/spaces/${params.name}/revision/${encodeURIComponent(params.revision ?? "HEAD")}?${search.toString()}`,
 		{
 			headers: {
 				...(accessToken ? { Authorization: `Bearer ${accessToken}` } : {}),
