ちくちく言葉チェッカーを作ってみた

2022/10/22に公開

はじめに

Google が公開している機械学習のフレームワークに TensorFlow というものがあります。
TensorFlow.js とは、JavaScript 製のフレームワークです。

今回は、Toxicity というモデルを使ってみました。
https://github.com/tensorflow/tfjs-models/tree/master/toxicity

作ったもの

  • こちらのサイトです。
  • テキストエリアに文章を入力して、ちくちく言葉であるかグラフで可視化します。
  • 現状英語のみ対応しています。

以下の環境で動作確認しています。

"@tensorflow-models/toxicity": "^1.2.2",
"@tensorflow/tfjs": "^3.20.0",
"chart.js": "^3.9.1",
"react": "^18.2.0",
"react-chartjs-2": "^4.3.1",
"react-dom": "^18.2.0",
"react-loading": "^2.0.3"

詳細

プロジェクト作成

今回は Vite + React + TypeScript の環境で開発しました

npm create vite@latest

https://vitejs.dev/guide/

@tensorflow/tfjs@tensorflow-models/toxicity をインストールします

$ npm install @tensorflow/tfjs @tensorflow-models/toxicity

使用例としては以下です

Toxicity.tsx
import React, { FC, useState } from 'react'
export const Toxicity: FC = () => {
  const threshold = 0.5
  const toxicityLabels = [
    'identity_attack',
    'insult',
    'obscene',
    'severe_toxicity',
    'sexual_explicit',
    'threat',
    'toxicity',
  ]

  toxicity.load(threshold, toxicityLabels).then((model) => {
    const sentences = ['ここにチェックしたい文章']
    model.classify(sentences).then((predictions) => {
      console.log(predictions)
    })
  })

  return (
    <div>
      //  略
    </div>
  )
}

console.log(predictions) の中身は以下のようになります。(数値はダミーです)

{
  "label": "identity_attack",
  "results": [{
    "probabilities": [0.9659664034843445, 0.03403361141681671],
    "match": false
  }]
},
{
  "label": "insult",
  "results": [{
    "probabilities": [0.08124706149101257, 0.9187529683113098],
    "match": true
  }]
},
...

上記の結果が返ってきた場合、入力した文章に対して以下のことが確認できます

  • identity_attack において、97%の確率で問題がなく、3%の確率で問題がある
  • insult において、8%の確率で問題がなく、92%の確率で問題がある
  • match は設定した閾値(threshold)を超えたかどうか

次にレスポンスの値をグラフで表示するために chart.jsreact-chartjs-2 をインストールします

npm install chart.js react-chartjs-2

https://react-chartjs-2.js.org/

各ラベルのデータを受け取り円グラフを作成します

PieChart.tsx
import { FC, memo } from 'react'
import { Chart as ChartJS, ArcElement, Tooltip, Legend } from 'chart.js'
import { Pie } from 'react-chartjs-2'
import type { Prediction } from '../../types/type'

ChartJS.register(ArcElement, Tooltip, Legend)

export const PieChart: FC<{ result: Prediction }> = memo(({ result }) => {
  const data = {
    labels: ['問題なし', '問題あり'],
    datasets: [
      {
        label: 'テスト',
        data: result.results[0].probabilities,
        backgroundColor: ['rgba(54, 162, 235, 0.2)', 'rgba(255, 99, 132, 0.2)'],
        borderColor: ['rgba(54, 162, 235, 1)', 'rgba(255, 99, 132, 1)'],
        borderWidth: 1,
      },
    ],
  }
  return (
    <div>
      <h2>{result.label}</h2>
      <Pie data={data} />
    </div>
  )
})

つまづいたこと

今回作成した Toxicity.tsx の親コンポーネントで、@tensorflow/tfjsをインポートしていると以下の issues にもあるとおりエラーになってしまいました。

https://github.com/tensorflow/tfjs/issues/2329

おわりに

  • TensorFlow.js と既存のモデルを使用することで機械学習にふれることができました。
  • 他にも面白そうなモデルがありました。
  • 今回のソースコードは以下で公開しています。

https://github.com/TakaShinoda/tfjs-react/tree/main/src/components/toxicity

GitHubで編集を提案

Discussion