WebGPU作为一项全新的图形API，带来了强大的硬件加速能力，彻底改变了前端开发的格局。过去只能在后台完成的高性能计算和图形渲染，现在可以在浏览器中直接实现，这意味着我们能够在浏览器中享受更快速、更高效的AI处理能力。对于AI模型来说，这不仅能够显著提高数据处理的速度，还意味着实时推理和计算不再仅仅依赖服务器，而是能够直接在用户的设备上完成。
而大模型，作为AI领域的近两年最热门的话题，不仅能够生成文字、图片，还能分析情感、做出决策。如果能够将WebGPU与大模型相结合，这无疑将带来更多创新和突破。无论是图像生成、自然语言处理，还是更复杂的AI推理应用，都能在端侧完成，大大降低对服务器资源的依赖。
是不是觉得这个想法很激动人心？
没错，咱们开源团队开源的白泽工具箱也完成了相关版本的开发，感兴趣的小伙伴可以来玩下。
安装步骤非常简单：下载并安装工具箱，然后通过百度网盘或Hugging Face获取所需的模型，放入指定文件夹后重启软件，便可以与AI进行对话了。
注意，第一次聊天的时候因为需要初始化模型，可能时间需要稍微多点（10来s，后续的聊天都是十分快的）
需要特别提醒的是，由于采用了WebGPU技术，这意味着不挑显卡，无论是AMD显卡还是NVIDIA显卡都能正常运行（理论上，国产摩尔线程显卡也可以支持）。不过，由于会消耗一定的显存，建议显卡至少为6GB，或者为集成显卡分配更多显存，以保证更好的性能。
快来体验吧，AI的世界，正等着你探索！
emmm，接下来，咱们就简单聊聊这个功能是怎么实现的吧。如果对实现不感兴趣的小伙伴可以移步啦。
这个功能模块咱们主要是依赖了transformers最近发布的transformers.js v3。通过它，咱们能够在浏览器中通过webgpu轻松加载和运行各种大规模AI模型，实现更加高效、快速的推理操作。
首先，咱们需要安装一些基础的依赖库，这些库将为模型的加载和推理提供支持
pnpm add @huggingface/transformers
新增一个worker文件，并利用@huggingface/transformers提供的AutoTokenizer和AutoModelForCausalLM构建AI对话模型的pipeline。
在worker中监听主线程中发送的相关的postmessage消息（监听generate事件），然后通过generate方法中的tokenizer.apply_chat_template调用大模型来进行对话。
ps:此处只粘贴了核心代码，详细可以去github中查看。
// worker.ts
class TextGenerationPipeline {
static model_id = null
static model = null
static tokenizer = null
static streamer = null
static async getInstance(progress_callback = null) {
// Choose the model based on whether fp16 is available
// this.model_id = "Phi-3-mini-4k-instruct";
this.model_id = "Phi-3-mini-4k-instruct"
this.tokenizer ??= AutoTokenizer.from_pretrained(this.model_id, {
legacy: true,
progress_callback,
})
this.model ??= AutoModelForCausalLM.from_pretrained(this.model_id, {
dtype: "q4",
device: "webgpu",
use_external_data_format: true,
progress_callback,
})
return Promise.all([this.tokenizer, this.model])
}
}
// 大模型生成
async function generate(messages) {
// Retrieve the text-generation pipeline.
const [tokenizer, model] = await TextGenerationPipeline.getInstance()
.then((res) => {
return res
})
.catch((e) => {
self.postMessage({
status: "errorMessage",
message: "模型加载失败，请到设置页面检查模型！",
})
return null
})
const inputs = tokenizer.apply_chat_template(messages, {
add_generation_prompt: true,
return_dict: true,
})
let allOutputText = ""
let startTime
let numTokens = 0
const cb = (output) => {
allOutputText += output
startTime ??= performance.now()
let tps
if (numTokens++ > 0) {
tps = (numTokens / (performance.now() - startTime)) * 1000
}
self.postMessage({
status: "update",
output: allOutputText.replace("<|end|>", ""),
tps,
numTokens,
})
}
// 监听主线程中的消息
self.addEventListener("message", async (e) => {
const { type, data } = e.data
switch (type) {
case "init":
env.remoteHost = data.remoteHost
console.log(env.remoteHost)
break
case "load":
load()
break
case "generate":
stopping_criteria.reset()
generate(data)
break
case "interrupt":
stopping_criteria.interrupt()
break
case "reset":
stopping_criteria.reset()
break
}
})
主进程中，主要就是构建messages聊天的消息队列，以及监听worker通过postmessage放回的流式聊天数据。
ps：此处只粘贴了核心代码（简化版），详细可以去github中查看。
export default function Chat() {
const workerRef = useRef<Worker | null>(null)
const sendMessage = (value: string) => {
const message: MessageItem = {
id: nanoid(8),
role: "user",
content: value,
}
workerRef.current?.postMessage({ type: "generate", data: [messages] })
}
useEffect(() => {
const worker = new Worker(new URL("./worker.ts", import.meta.url), {
type: "module",
})
workerRef.current = worker
const onMessageReceived = (e) => {
switch (e.data.status) {
case "errorMessage":
message.error(e.data.message)
setIsChat(false)
break
case "start":
break
case "update":
const { output, tps, numTokens } = e.data
console.log(output, tps, numTokens)
const newList = cloneDeep(messageListRef.current)
const lastMessage = newList[newList.length - 1]
if (lastMessage.role === "user") {
const newMessage = {
id: nanoid(8),
role: "assistant",
content: output,
}
changeMessageList([...newList, newMessage])
} else {
lastMessage.content = output
changeMessageList(newList)
}
break
case "complete":
setIsChat(false)
console.log("Generation complete")
break
}
}
workerRef.current.addEventListener("message", onMessageReceived)
return () => {
workerRef.current?.removeEventListener("message", onMessageReceived)
}
}, [])
ok，到这里，基本上核心代码都已经写完了，但是乍一看好像也不需要用到端侧啊？为啥你们需要上端侧呢？直接网页不可以吗？
确实，直接在网页上运行也是可以的。但是这其中有一个非常现实的问题需要考虑。webgpu中使用的模型通常都是经过量化处理的，尽管经过量化，模型的大小依然可能达到3GB+。这样子对咱们的网络压力是十分大的。而且加载一个模型可能得先等几分钟，这对用户的体验是极度不友好的。因此我想，这也是webgpu版本的大模型出来了，但是好像还没什么人使用上。
但是如果咱们使用的是端侧，那就不一样了，端侧咱们能做的事情就变多了。像在electron中，咱们就可以直接通过file协议来加载本地的文件，这样子初始化的时间就会降低到10来s(这应该属于一个可接受的事件，小伙伴们可以通过demo来试试)
ps:electron的最新版本好像不再支持file协议直接加载文件，不过31.7.5是可以的。
注意如果transformers.js想使用本地模型，则需要添加如下配置
// worker.ts
env.remotePathTemplate = "{model}"
env.useBrowserCache = true
env.remoteHost = "你的host"
整体的架构图如下：
小结
关于端侧AI实战的介绍就到这里啦。那么，小伙伴们觉得WebGPU+大模型的结合，会不会改变前端生态呢？
小羽觉得，端侧的大模型确实存在一定的需求。首先，它可以充分利用用户端的显卡，降低对服务器的依赖，减轻服务器的负担。此外，端侧模型还支持在完全断网的环境下工作，特别适合那些网络状况较差的用户，确保无论在何种网络环境下都能顺畅运行。更重要的是，对于公司来说，某些敏感数据无法暴露到云端时，端侧大模型也能有效地进行本地化处理和分析，确保数据安全。而webgpu+electron则可以降低适配的成本，简直完美~
总之，WebGPU和端侧大模型的结合，既能提升用户体验，又能开辟更多创新应用的空间，对前端生态的未来将产生无限的想象空间。
加入我们
白泽大舞台，有胆你就来！
