🦔
Vercel AI SDK Core + Gemma 2を使って記事のカテゴリ分類タスクを開発する
Vercel AI SDKはLLMを活用したアプリケーションを構築できるように設計されたTypeScript製のツールキットです。チャットボットの構築に伴う定型的なコードを排除して、本質的な部分の開発に集中することを可能にします。
しかしVercel AI SDKは何もチャットボット開発だけに使えるわけではありません。たとえばLangChainのように様々なモデルとツールを組み合わせて、複雑なタスクを解決するデータパイプラインを構築することも可能です（Pythonの豊富なライブラリと比べるとまだまだ実現できることは限定的ですが）。
Vercel AI SDKは、主に次の3つのコンポーネントで構成されています。
- AI SDK Core: LLMを使用してテキスト、構造化オブジェクト、ツール呼び出しを生成するための統一APIを提供します。
- AI SDK UI: チャットや生成UIを迅速に構築するための一連のフレームワークに依存しないフックを提供します。
- AI SDK RSC: Reactサーバーコンポーネント（RSC）を使用して生成UIをストリーミングするためのライブラリです。
このうちUIを構築するコンポーネントは「AI SDK UI」と呼ばれ、「AI SDK Core」はモデルの呼び出しやデータの加工を行うためのライブラリです。
チャットボット開発以外のタスクにもVercel AI SDKを活用できることを示す一例として、この記事ではAI SDK Coreのみを使ってSpeakerDeckの投稿タイトルからカテゴリを判定するタスクを開発します。具体的には、はてなブックマークの検索結果からSpeakerDeckの投稿データを取得し、AI SDK Coreを用いてカテゴリ分類と評価を行う一連のNode.jsプログラムを作成します。なぜSpeakerDeckかというと単に私がスライドをまとめ読みしたいからです。
類似の取り組みとしてはメルカリエンジニアリングに投稿された「LLMを活用した大規模商品カテゴリ分類への取り組み」での「商品データの正解カテゴリ予測」が挙げられます。
記事の著者のML_Bearさんはコスト効率と処理速度のバランスを踏まえChatGPT 3.5 turboを採用していますが、最近のローカルLLMの進化は目覚ましいものがあります。その中でもGemma 2は、Googleが開発した最新の言語モデルであり、性能で高い評価を受けています。
なのでこの記事では、Gemma 2（9Bモデル）を使いつつ、AI SDK Coreを使ってカテゴリ分類タスクを開発してみます。ローカル環境での実行を前提としているのでAPI呼び出しコストがかからないのも魅力です。
データの取得と前処理
はてなブックマークの検索結果から、最新登録されたSpeakerDeckの投稿データを取得します。これには、ブラウザ上でスクレイピングを行います。
- はてなブックマークの検索結果ページにアクセスします。
- ブラウザのコンソールを開き、以下のスクリプト(scraping.js)をコピペして実行します。
- 以下のようなCSV ファイルがダウンロードされます。
title,url,username,bookmarkCount,description,date,tags
"1","健康第一！MetricKitで始めるアプリの健康診断 / App Health Checkups Starting with MetricKit","https://speakerdeck.com/nekowen/app-health-checkups-starting-with-metrickit","nekowen","3","iOSDC Japan 2024 Day0 Track B 18:10 -","2024/09/01 20:50",""
"2","LR で JSON パーサーを作る / Coding LR JSON Parser","https://speakerdeck.com/junk0612/coding-json-lr-parser","junk0612","3","大阪Ruby会議04でのスポンサーLTです。 番組内で作成したパーサーはこちら → https://gist.github.com/junk0612/c6bc79776724ab8de1857b5d1fc1b360","2024/09/01 15:26",""
"3","Rubyとクリエイティブコーディングの輪の広がり / The Growing Circle of Ruby and Creative Coding","https://speakerdeck.com/chobishiba/the-growing-circle-of-ruby-and-creative-coding","chobishiba","3","RubyKaigi 2024 follow up 例のあれ、どうなりました？ で発表した資料です RubyKaigi 2024でのLT(元ネタ) https://speakerdeck.com/chobishiba/enjoy-creative-coding-with-ruby-rubykaig…","2024/09/01 07:34",""
scraping.js
/**
* 1. Visit https://b.hatena.ne.jp/site/speakerdeck.com/?page=1
* 2. Open the browser console (F12)
* 3. Copy and paste the following code
*/
const headers = ['title', 'url', 'username', 'bookmarkCount', 'description', 'date', 'tags'];
let csvContent = headers.join(',') + '\n';
const entries = Array.from(document.querySelectorAll('.entrylist-image-entry')).slice(1); // To exclude the first entry which is an ad
entries.forEach(entry => {
const titleElement = entry.querySelector('.entrylist-contents-title a');
const title = titleElement.textContent.trim();
const urlElement = entry.querySelector('.entrylist-contents-title a');
const url = urlElement ? urlElement.href : '';
const username = url.split('/')[3];
const bookmarkCountElement = entry.querySelector('.entrylist-contents-users span');
const bookmarkCount = bookmarkCountElement ? bookmarkCountElement.textContent.trim() : '';
const descriptionElement = entry.querySelector('.entrylist-contents-description');
const description = descriptionElement ? descriptionElement.textContent.trim() : '';
const dateElement = entry.querySelector('.entrylist-contents-date');
const date = dateElement ? dateElement.textContent : '';
const tagsElements = entry.querySelectorAll('.entrylist-contents-tags a');
const tags = tagsElements.length > 0 ? Array.from(tagsElements).map(a => a.textContent).join(', ') : '';
const id = entries.indexOf(entry) + 1; // articale_id
csvContent += [id, title, url, username, bookmarkCount, description, date, tags].map(value => `"${String(value).replace(/"/g, '""')}"`).join(',') + '\n';
});
console.log(csvContent);
const IS_DOWNLOAD_ENABLED = true;
if (IS_DOWNLOAD_ENABLED) {
const encodedUri = encodeURI('data:text/csv;charset=utf-8,' + csvContent);
const link = document.createElement('a');
link.setAttribute('href', encodedUri);
link.setAttribute('download', 'data.csv');
document.body.appendChild(link); // Required for FF
link.click();
document.body.removeChild(link); // Clean up
}
Vercel AI SDK Core を用いたモデルの呼び出し
データが用意できたので次はNode.jsプロジェクトを作成して、Vercel AI SDKによる分類タスクを実装していきます。
プロジェクトを作成し、依存するパッケージをインストールします。
npm init -y
npm install @ai-sdk/core @ai-sdk/openai zod ollama-ai-provider
package.json は以下のように"type": "module"を追加してください。
package.json
{
"name": "speakerdeck-classification",
"version": "1.0.0",
"description": "",
"main": "index.js",
+ "type": "module",
"scripts": {
"start": "node index.js"
},
// ...
}
SpeakerDeck の投稿タイトルを入力として、Ollamaを通してGemma 2に推論させ、カテゴリを予測します。
index.js
import { generateObject } from "ai"
import { openai } from "@ai-sdk/openai"
import fs from 'fs';
import { createOllama } from "ollama-ai-provider"
import { z } from "zod"
const ollama = createOllama()
/*
* INSTALLATION: 下記の全記事リストからカテゴリーとサブカテゴリーを選択してください。もし記事がどのカテゴリーやサブカテゴリーにも当てはまらない場合は、「others」を選択してください。
*/
const INSTALLATION = 'Select category and subcategory from the folowwing list of all articles. if the article does not fit any of the categories or subcategory, select "others".'
const data = fs.readFileSync('data.csv', 'utf8');
const lines = data.split('\n');
const csvContent = lines.slice(0, 10); // 先頭10行を抽出
const result = await generateObject({
model: ollama("gemma2"),
temperature: 0.0,
schema: z.object({
result: z.array(
z.object({
id: z.number(),
title: z.string(),
category: z.string(),
subcategory: z.string(),
})
),
}),
prompt: `${INSTALLATION}
CATEGORIES:
${CATEGORIES}
ARTICLES:
${csvContent}`,
};
generateObjectは、指定されたZodスキーマに一致する型指定された構造化オブジェクトを生成します。これは結果をさらにプログラムで加工する際に便利です。
CATEGORIES定数には、以下のカテゴリーとサブカテゴリーをMarkdown形式で記述しました。ここは適宜変更してください。
const CATEGORIES = `
CATEGORIES:
- frontend
- react
- vue
- backend
- java
- python
- database
- ruby
- php
- golang
- rust
- mobile
- ios
- android
- flutter
- game
- unity
- unreal-engine
- cloud
- aws
- azure
- gcp
- ai
- llm
- ml
- management
- engineering
- product
- design
- ui
- css
`;
ARTICLES以下には、先ほど取得したCSVデータを埋め込みます。試したところGemma 2はトークン数が増えるほど構造化に失敗しやすかったため、ここでは先頭10行のみを抽出しています。
このコードを実行するためには、ローカル環境でollama serverを起動しておく必要があります。
❯ ollama pull gemma2
❯ ollama start
2024/09/09 17:26:00 routes.go:1125: INFO server config env="map[OLLAMA_DEBUG:false OLLAMA_FLASH_ATTENTION:false OLLAMA_HOST:http://127.0.0.1:11434 OLLAMA_KEEP_ALIVE:5m0s OLLAMA_LLM_LIBRARY: OLLAMA_MAX_LOADED_MODELS:0 OLLAMA_MAX_QUEUE:512 OLLAMA_MODELS:/Users/kstg/.ollama/models OLLAMA_NOHISTORY:false OLLAMA_NOPRUNE:false OLLAMA_NUM_PARALLEL:0 OLLAMA_ORIGINS:[http://localhost https://localhost http://localhost:* https://localhost:* http://127.0.0.1 https://127.0.0.1 http://127.0.0.1:* https://127.0.0.1:* http://0.0.0.0 https://0.0.0.0 http://0.0.0.0:* https://0.0.0.0:* app://* file://* tauri://*] OLLAMA_RUNNERS_DIR: OLLAMA_SCHED_SPREAD:false OLLAMA_TMPDIR:]"
❯ node index.js
結果をconsole.logで出力してみましょう。
console.log(result.object.result);
❯ node index.js
ID: 1, Title: 健康第一！MetricKitで始めるアプリの健康診断 / App Health Checkups Starting with MetricKit, Category: mobile, Subcategory: ios
index.js:101
ID: 2, Title: LR で JSON パーサーを作る / Coding LR JSON Parser, Category: backend, Subcategory: ruby
index.js:101
ID: 3, Title: Rubyとクリエイティブコーディングの輪の広がり / The Growing Circle of Ruby and Creative Coding, Category: backend, Subcategory: ruby
分類結果の確認
予測されたカテゴリと実際のスライドを比較し、適切に分類されたかを確認します。
ところでこれを自動化してみましょう。Gemma 2の予測結果をgpt-4oに評価させて、間違っていそうな場合は適切なカテゴリを提案してもらいます。
index.js
const report = result.object.result.map(({ id, title, category, subcategory }) => ({
id,
title,
category,
subcategory
}));
const csvReport = [
'id, title, category, subcategory',
...report.map(({ id, title, category, subcategory }) => `${id}, ${title}, ${category}, ${subcategory}`)
].join('\n');
const evaluation = await generateObject({
model: openai("gpt-4o"),
temperature: 0.0,
schema: z.object({
results: z.array(
z.object({
id: z.number(),
correct: z.boolean(),
suggestion: z.string().optional(),
})
),
}),
prompt: `以下の記事のカテゴリとサブカテゴリの分類結果を評価してください。「リストにないカテゴリ」や「カテゴリが適切でない」等の誤った分類があれば修正してください。
INSTALLATION:
${INSTALLATION}
CATEGORIES:
${CATEGORIES}
${csvReport}`,
});
const resultsWithTitles = result.object.result.map((item, index) => ({
...evaluation.object.results[index],
title: item.title,
category: item.category,
subcategory: item.subcategory,
}));
for (const { id, correct, title, category, subcategory, suggestion } of resultsWithTitles) {
if (!correct) {
console.log(`ID:${id} 「${title}」は${category}, ${subcategory} として分類されましたが、${suggestion} が適切な分類です。`);
}
}
以下のように出力されました。
❯ node index.js
ID:5 「実践的なバグバウンティ入門」はbackend, security として分類されましたが、others が適切な分類です。
ID:6 「ECMAScript、Web標準の型はどう管理されているか / How ECMAScript and Web standards types are maintained」はfrontend, javascript として分類されましたが、others が適切な分類です。
このように複数のモデルを組み合わせて、タスクによって最適なモデルを横断的に選択することができます。同じインターフェイスでパラメータ調整やデータ加工を行うことができるため、お手軽にチューニングを行うことができます。
まとめ
この記事では、Vercel AI SDK Coreを用いてSpeakerDeckの投稿タイトルからカテゴリを判定するタスクを開発しました。
Vercelのプラットフォームはフロントエンドのデプロイメントサービスとして有名ですが、昨今のAIブームに乗り、AI SDKやv0.devなどのAI関連のサービスに力を入れているようです。CEOのXのポストも半数以上がAI関連です（これもAI SDKで分析しました）。まさに社運をかけたプロジェクトと言えるでしょう。Vercel儲かっててサ終しないくれ。
おまけ：DuckDB便利
説明のためにDB環境の説明を省略しましたが、著者はデータ保存にDuckDBを使っています。DuckDBは、PythonのPandasのようなデータフレームをSQLで操作できるローカルデータベースです。
DuckDBを使うと簡単にデータを扱うことができます。CSVファイルを取り込むには以下のようにします。
// npm install duckdb
import duckdb from 'duckdb';
export function setupDatabase() {
const db = new duckdb.Database('articles.db');
const conn = db.connect();
conn.run(`
CREATE OR REPLACE SEQUENCE articles_seq;
CREATE TABLE IF NOT EXISTS articles (
article_id INTEGER PRIMARY KEY DEFAULT nextval('articles_seq'),
title VARCHAR,
url VARCHAR UNIQUE,
username VARCHAR,
bookmarkCount INTEGER,
description VARCHAR,
date DATE,
tags VARCHAR
);
CREATE TABLE IF NOT EXISTS classifications (
article_id INTEGER UNIQUE,
category VARCHAR,
subcategory VARCHAR,
FOREIGN KEY (article_id) REFERENCES articles(article_id)
);
`);
return { db, conn };
}
❯ duckdb articles.db -s "COPY (
SELECT nextval('articles_seq') as article_id, *
FROM read_csv('import.csv')
) TO 'articles';"
❯ duckdb articles.db -s "SELECT * FROM articles LIMIT 3;"
┌────────────┬──────────────────────┬──────────────────────┬────────────┬───────────────┬─────────────────────────────────────────────────────────────────────┬────────────┬───────────────────────────────────────┐
│ article_id │ title │ url │ username │ bookmarkCount │ description │ date │ tags │
│ int32 │ varchar │ varchar │ varchar │ int32 │ varchar │ date │ varchar │
├────────────┼──────────────────────┼──────────────────────┼────────────┼───────────────┼─────────────────────────────────────────────────────────────────────┼────────────┼───────────────────────────────────────┤
│ 1 │ 生成AIの二大潮流と… │ https://speakerdec… │ koukyo1994 │ 97 │ https://yans.anlp.jp/entry/yans2024 での講演スライドです。 │ 2024-09-07 │ AI, あとで読む, IT, 仕事, *あとで読む │
│ 2 │ なぜクラウドサービ… │ https://speakerdec… │ shuta13 │ 43 │ Web Developer Conference 2024 にて行ったセッションの資料です。 │ 2024-09-07 │ サービス, あとで読む │
│ 3 │ React Aria で実現… │ https://speakerdec… │ ryo_manba │ 22 │ Web Developer Conference 2024 (2024/09/07) での発表資料 https://w… │ 2024-09-07 │ react, あとで読む │
└────────────┴──────────────────────┴──────────────────────┴────────────┴───────────────┴─────────────────────────────────────────────────────────────────────┴────────────┴───────────────────────────────────────┘
import { setupDatabase } from './db.js';
const { db, conn } = setupDatabase();
db.all(`SELECT * FROM articles LIMIT 10`, (err, records) => {
if (err) {
console.error(err);
return;
}
// ...
const stm = conn.prepare('INSERT INTO classifications (category, subcategory, article_id) VALUES (?, ?, ?) ON CONFLICT DO UPDATE SET category = EXCLUDED.category, subcategory = EXCLUDED.subcategory;');
for (const { id, title, category, subcategory } of result.object.result) {
console.log(`Updating ID: ${id} title: ${title}, category: ${category}, subcategory: ${subcategory}`);
stm.run(category, subcategory, id);
}
stm.finalize();
db.all(`SELECT title, category, subcategory, url from articles inner join classifications on classifications.article_id = articles.article_id where classifications.category = 'others' order by category`, (err, rs) => {
if (err) {
console.error(err);
return;
}
for (const r of rs) {
console.log(`Title: ${r.title}, Category: ${r.category}, Subcategory: ${r.subcategory}, URL: ${r.url}`);
}
});
});
});
❯ node index.js
Title: EitherT_with_Future, Category: others, Subcategory: null, URL: https://speakerdeck.com/aoiroaoino/eithert-with-future
Title: null or undefined, Category: others, Subcategory: others, URL: https://speakerdeck.com/susisu/null-or-undefined
Title: Datadogマニアック機能活用術, Category: others, Subcategory: null, URL: https://speakerdeck.com/biwashi/mastering-datadogs-advanced-features
Discussion