Fix generate script after new snippets

huggingface · Nov 15, 2024 · 228b216 · 228b216
1 parent 23a1d86
commit 228b216
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 23 deletions.
diff --git a/scripts/api-inference/pnpm-lock.yaml b/scripts/api-inference/pnpm-lock.yaml
diff --git a/scripts/api-inference/scripts/generate.ts b/scripts/api-inference/scripts/generate.ts
@@ -115,19 +115,22 @@ export function getInferenceSnippet(
   id: string,
   pipeline_tag: PipelineType,
   language: InferenceSnippetLanguage,
+  config?: JsonObject,
+  tags?: string[],
 ): string | undefined {
   const modelData = {
     id,
     pipeline_tag,
     mask_token: "[MASK]",
     library_name: "",
-    config: {},
-    tags: [],
+    config: config ?? {},
+    tags: tags ?? [],
   };
   // @ts-ignore
   if (HAS_SNIPPET_FN[language](modelData)) {
     // @ts-ignore
-    return GET_SNIPPET_FN[language](modelData, "hf_***");
+    const snippets = GET_SNIPPET_FN[language](modelData, "hf_***");
+    return Array.isArray(snippets) ? snippets[0].content : snippets.content;
   }
 }
 
@@ -472,25 +475,13 @@ function fetchChatCompletion() {
     );
 
     const mainModel = DATA.models[task.name][0];
-    const mainModelData = {
-      // @ts-ignore
-      id: mainModel.id,
-      pipeline_tag: task.pipelineTag,
-      mask_token: "",
-      library_name: "",
-      // @ts-ignore
-      tags: ["conversational"],
-      // @ts-ignore
-      config: mainModel.config,
-    };
-
     const taskSnippets = {
       // @ts-ignore
-      curl: GET_SNIPPET_FN["curl"](mainModelData, "hf_***"),
+      curl: getInferenceSnippet(mainModel.id, task.pipelineTag, "curl", mainModel.config, ["conversational"]),
       // @ts-ignore
-      python: GET_SNIPPET_FN["python"](mainModelData, "hf_***"),
+      python: getInferenceSnippet(mainModel.id, task.pipelineTag, "python", mainModel.config, ["conversational"]),
       // @ts-ignore
-      javascript: GET_SNIPPET_FN["js"](mainModelData, "hf_***"),
+      javascript: getInferenceSnippet(mainModel.id, task.pipelineTag, "js", mainModel.config, ["conversational"]),
     };
     DATA.snippets[task.name] = SNIPPETS_TEMPLATE({
       taskSnippets,