Merge pull request #323 from Linchenn/main

Add LLM Inference sample for Web
google-ai-edge · Mar 6, 2024 · 76ee917 · 76ee917
2 parents d379c85 + 66ed93b
commit 76ee917
Show file tree

Hide file tree

Showing 3 changed files with 114 additions and 0 deletions.
diff --git a/examples/llm_inference/js/README.md b/examples/llm_inference/js/README.md
@@ -0,0 +1,19 @@
+# MediaPipe LLM Inference task for web
+
+## Overview
+
+This web sample demonstrates how to use the LLM Inference API to run common text-to-text generation tasks like information retrieval, email drafting, and document summarization, on web.
+
+
+## Prerequisites
+
+* A browser with WebGPU support (such as Chrome, Edge or Opera).
+
+## Running the demo
+
+Follow the following instructions to run the sample on your device:
+1. Make a folder for the task, named as `llm_task`, and copy the [index.html](https://github.com/googlesamples/mediapipe/blob/main/examples/llm_inference/js/index.html) and [index.js](https://github.com/googlesamples/mediapipe/blob/main/examples/llm_inference/js/index.js) files into your `llm_task` folder.
+2. Download one of the [compatible models](https://developers.google.com/mediapipe/solutions/genai/llm_inference#models) that you want to run, into the `llm_task` folder.
+3. In your `index.js` file, update []`modelFileName`](https://github.com/googlesamples/mediapipe/blob/main/examples/llm_inference/js/index.js#L23) with your model file's name.
+4. Run `python3 -m http.server 8000` under the `llm_task` folder to host the three files.
+5. Open `localhost:8000` in Chrome. Then the button on the webpage will be enabled when the task is ready (~10 seconds).
diff --git a/examples/llm_inference/js/index.html b/examples/llm_inference/js/index.html
@@ -0,0 +1,30 @@
+<!-- Copyright 2024 The MediaPipe Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. -->
+
+<!doctype html>
+<html lang="en">
+<head>
+    <title>LLM Inference Web Demo</title>
+</head>
+<body>
+    Input:<br />
+    <textarea id="input" style="height: 300px; width: 600px"></textarea><br />
+    <input type="button" id="submit" value="Get Response" disabled /><br />
+    <br />
+    Result:<br />
+    <textarea id="output" style="height: 300px; width: 600px"></textarea>
+    <script src="https://cdn.jsdelivr.net/npm/@mediapipe/tasks-genai/genai_bundle.js"></script>
+    <script type="module" src="index.js"></script>
+</body>
+</html>
diff --git a/examples/llm_inference/js/index.js b/examples/llm_inference/js/index.js
@@ -0,0 +1,65 @@
+// Copyright 2024 The MediaPipe Authors.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//      http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// ---------------------------------------------------------------------------------------- //
+
+import {FilesetResolver, LlmInference} from 'https://cdn.jsdelivr.net/npm/@mediapipe/tasks-genai';
+
+const input = document.getElementById('input');
+const output = document.getElementById('output');
+const submit = document.getElementById('submit');
+
+const modelFileName = 'llm.tflite'; /* Update the file name */
+
+/**
+ * Display tokens to the output text box.
+ */
+function displayNewTokens(tokens, complete) {
+  output.textContent += tokens;
+
+  if (complete) {
+    if (!output.textContent) {
+      output.textContent = 'Result is empty';
+    }
+    submit.disabled = false;
+  }
+}
+
+/**
+ * Main function to run LLM Inference.
+ */
+async function runDemo() {
+  const genaiFileset = await FilesetResolver.forGenAiTasks(
+      'https://cdn.jsdelivr.net/npm/@mediapipe/tasks-genai/wasm');
+  let llmInference;
+
+  submit.onclick = () => {
+    output.textContent = '';
+    submit.disabled = true;
+    llmInference.generateResponse(input.value, displayNewTokens);
+  };
+
+  submit.value = 'Loading the model...'
+  LlmInference
+      .createFromModelPath(genaiFileset, modelFileName)
+      .then(llm => {
+        llmInference = llm;
+        submit.disabled = false;
+        submit.value = 'Get Response'
+      }).catch(() =>{
+        alert('Failed to initialize the task.');
+      });
+}
+
+runDemo();