Neovimで手動のスピードで入力しよう！ - エムスリーテックブログ

AI・機械学習チームの北川(@kitagry)です。最近Neovimの設定ファイルをinit.vimからinit.luaに移動させました。 Lua化したからには何かそれっぽいことをしたいなと思い、機械学習しつつ何かアプリケーションっぽいものを作成することにしました。

タイトルは以下の本のオマージュです。神本なので、おすすめです。

www.kadokawa.co.jp

Neovimで機械学習をする環境を用意する
- NeovimからLuaファイルを呼び出す仕組み
- torchをインストールする
Neovimでmnistを学習する
Neovimで文字を書く
手動のスピードで入力する
まとめ
- おまけ
We are hiring!!

Neovimで機械学習をする環境を用意する

まずはじめにNeovimで機械学習を扱う仕組みを用意する必要があります。今回はLuaのライブラリであるtorchを利用します。*1 最近ではPyTorchが有名だと思いますが、torchはその前身であるライブラリでLuaで利用可能な機械学習ライブラリです。

しかし、torchのNeovimプラグインは存在しないため自身でtorchを利用する環境を用意する必要があります。そのために、どのようにすればNeovimプラグインではないLuaライブラリをNeovimから呼び出すことが出来るかを知る必要があります。

NeovimからLuaファイルを呼び出す仕組み

Vimではどのディレクトリから設定ファイルやプラグインファイルを読み込むかをruntimepathというオプションで決定しています。 NeovimのLuaファイルはruntimepathの各ディレクトリにある /lua/?.lua もしくは /lua/?/init.lua というファイル形式のファイルを読み込むことが出来ます。 (詳しくは :help lua-package-path)

そこでNeovimからtorchを呼び出すためには以下のようなディレクトリ構成を用意し、そのディレクトリをruntimepathに含めれば良いです。

.
└── lua
    └── torch.lua

上のような仕組みにすることによってNeovim内で :lua require('torch') でプラグインとして呼び出すことができます。

torchをインストールする

torchは以下のようにインストールします. http://torch.ch/docs/getting-started.html

git clone https://github.com/torch/distro.git ~/torch --recursive
cd ~/torch; bash install-deps
./install.sh

これを行うと ./install/share/lua/5.1/* にtorchを呼び出すためのLuaファイルが作成されます。このディレクトリ名(5.1)だとNeovimから呼び出せないため、Luaという名前になるようにする必要があります。

$ cp -r ./install/share/lua/5.1/* ./lua

次にこのディレクトリをruntimepathに加えます。

# vimの場合
echo "set runtimepath+=$(pwd)" >> ~/.config/nvim/init.vim

# Luaの場合
echo "vim.api.nvim_command('set runtimepath^=$(pwd)')" >> ~/.config/nvim/init.lua

これで、Neovimからtorchが使えるようになりました。試しにNeovim上でtorchを呼び出してみます。

:lua torch = require('torch')
:lua print(vim.inspect(torch))
{
  Allocator = <1>{
    <metatable> = <2>{
      __add = <function 1>,
...

無事にtorchをインストール出来ていることがわかりました。

Neovimでmnistを学習する

torchを呼び出すことに成功したので、次にmnistの学習をします。学習の全体のコードについてはdotfilesのマージされることはないであろうPRをご覧ください。

学習用のmnistのデータはhttps://github.com/torch/tutorialsのA_datasets/mnist.luaを実行して生成します。このファイルを動かすためにLuaの実行環境が必要になりますが、先ほどの手順でNeovim上でtorchを実行する環境が整っていれば以下のように動かすことが出来ます。

$ nvim -l mnist.lua

または、Neovimを起動時にLuaファイルを実行するには以下のように動かせます。

:luafile mnist.lua

データが用意できたので学習用のコードを書いていきます。学習は簡単に以下のようなモデルで行います。

local nn = require('nn')

local model = nn.Sequential()
model:add(nn.View(32*32))
model:add(nn.Linear(32*32, 1000)):add(nn.ReLU())
model:add(nn.Linear(1000, 1000)):add(nn.ReLU())
model:add(nn.Linear(1000, 10))
model:add(nn.LogSoftMax())
model = model:float()

このモデルをNeovim上で学習させます。これ以降は特に工夫したとかはなく、無骨に実装をおこなっただけなのでPRをご覧ください。

Neovim上で学習をしている様子(この間編集ができないのでエディタってなんだっけなと思考を巡らせることが出来ます)

Neovimで文字を書く

ただ、mnistの判定ができるモデルがNeovim上で動くだけでは面白くないので、Neovim上で文字を書くようにします。

Neovim上で文字を書くコードには notomo/gesture.nvim のコードを一部改変して利用しました。このプラグインはNeovim上でマウスジェスチャーを使うことによって画面の分割などを行うプラグインです。正直このプラグインの存在を知らなかったら、この記事は世の中に出なかったと思います。正直どうやって動いてるのかよく分からん。。筆者が書かれている記事によると、透明なフローティングウィンドウを用意し、その上にハイライトすることによって描画をおこなっているようです。

このプラグインをforkし、描画の各点の場所を返すように変更したコードを作成しました。以下のようにすることによって各点の場所を取得できます。

vim.keymap.set("n", "<LeftRelease>", function()
    local points = require("gesture").finish()
    print(vim.inspect(points))
end)

手動のスピードで入力する

さて、文字を書いて各点を取得する所までは作成できました。ここから先ほど学習したモデルで推論する用のデータを入れるためのコードを書きます。

文字サイズは32*32で、0~255の値を持つのでいい感じにデータを整形しつつ作成します。あとはモデルで推論した結果を現在の行の最後に追加していきます。

-- 文字を推測
-- [0.001, 0.001, 1, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001]のような形で出てくる
local output = model:forward(image)

-- 最も高い数値のindexを取得
local function argmax_1D(v)
  local length = v:size(1)
  assert(length > 0)

  -- examine on average half the entries
  local maxValue = torch.max(v)
  for i = 1, v:size(1) do
    if v[i] == maxValue then
      return i
    end
  end
end

-- カーソルがある行に推測した数値を追加する
local current_line = vim.api.nvim_get_current_line()
vim.cmd("call setline('.', '" .. current_line .. tostring(argmax_1D(output)-1) .. "')")