From 04e6de6f0a70703ca4f6aa489d2858a288e803d2 Mon Sep 17 00:00:00 2001
From: "Yuichiro Tachibana (Tsuchiya)" <t.yic.yt@gmail.com>
Date: Tue, 24 Oct 2023 18:00:21 +0900
Subject: [PATCH] Write a test to use pipeline with a
 zero-shot-image-classification model

---
 pyodide-e2e/src/tests/pipeline.test.ts | 41 ++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 pyodide-e2e/src/tests/pipeline.test.ts

diff --git a/pyodide-e2e/src/tests/pipeline.test.ts b/pyodide-e2e/src/tests/pipeline.test.ts
new file mode 100644
index 0000000..1f00aa2
--- /dev/null
+++ b/pyodide-e2e/src/tests/pipeline.test.ts
@@ -0,0 +1,41 @@
+import type { PyodideInterface } from "pyodide";
+import { beforeEach, describe, it, expect } from "vitest";
+import { setupPyodideForTest } from "./utils";
+
+describe("transformers.pipeline", () => {
+  let pyodide: PyodideInterface;
+
+  beforeEach(async () => {
+    pyodide = await setupPyodideForTest();
+  });
+
+  it("works with zero-shot-image-classification model", async () => {
+    await fetch("https://huggingface.co/spaces/gradio/image_mod/resolve/main/images/lion.jpg")
+      .then((response) => response.blob())
+      .then((blob) => blob.arrayBuffer())
+      .then((arrayBuffer) => {
+        const fileData = new Uint8Array(arrayBuffer);
+        const filePath = "/tmp/cheetah.jpg";
+        pyodide.FS.writeFile(filePath, fileData);
+      });
+
+    await pyodide.runPythonAsync(`
+from transformers_js import import_transformers_js, as_url
+
+transformers = await import_transformers_js()
+pipeline = transformers.pipeline
+pipe = await pipeline('zero-shot-image-classification')
+
+image = "/tmp/cheetah.jpg"
+
+data = await pipe(as_url(image), ["tower", "lion", "flower"])
+result = {item['label']: round(item['score'], 2) for item in data}
+`);
+    const resultMap = await pyodide.globals.get("result").toJs();  // Python's dict to JS's Map
+    const resultObj = Object.fromEntries(resultMap);
+    expect(Object.keys(resultObj)).toEqual(["tower", "lion", "flower"]);
+
+    const topLabel = Object.keys(resultObj).reduce((a, b) => resultObj[a] > resultObj[b] ? a : b);
+    expect(topLabel).toEqual("lion");
+  });
+});