🔎

日本語でコードを検索したい(on Neovim)

に公開
Neovim
treesitter
embedding
RAG
tech

作ったもの

日本語や英語でコードを検索するコマンドを実装して、Neovimに統合してみました。

「拡張子を調べる関数」と入力すると、拡張子やファイル形式に関連した関数が表示されて、「ファイルを読み取ってチャンキングする関数」と入力するとchunks変数やreadFileAndChunking関数などが候補に入っています。

デモ目的なので、検索結果からジャンプしたりする機能は搭載していません。純粋な検索のみです（ジャンプなどの実装自体はそれほど難しくはないと思います）。

再現手順

プラグイン自体がchopgrepに依存しているため、これを導入する必要があります。
https://github.com/sirasagi62/chopgrep
デモ目的+開発中+自分が使うことだけを前提にしているので、bun以外での動作/npmでのインストールに対応していません。よって、bunが動作することが必須です。

$ git clone https://github.com/sirasagi62/chopgrep
$ bun install
$ bun link

で、chopgrepコマンドが使えるようになります。

https://github.com/sirasagi62/chopgrep.nvim

chopgrep.nvimは通常のパッケージマネージャーで導入できます。Jetpackであれば

require('jetpack.packer').add {
-- other plugins...
  { 'sirasagi62/chopgrep.nvim',
    requires = {
      'MunifTanjim/nui.nvim',
      'grapp-dev/nui-components.nvim'
    }
  },
}

で導入できるはずです。

使い方

chopgrepは

  1. コードをチャンクに分割してベクトル埋め込みとしてSqliteにいれる機能
  2. Sqliteからコードをベクトル検索する機能
    の2つを持っています。
    プロジェクトのルートディレクトリでシェルを使って
$ chopgrep index # シェル

もしくはNeovimで

:ChopgrepIndex

とすることで、code_chunks.dbを生成します。
あとは

$ chopgrep query "<Query>"

とすることでコードを検索できます。READMEにあるように、-jオプションを加えることでjson出力できる他、表示件数も変更できます。

Neovimでは:Chopgrepとすることで冒頭の対話的なUIが立ち上がります。Neovim版では、Tabで検索窓、リスト、プレビューを移動できます。リストでEnterキーを押すことでPreviewを確認することができます。コードにジャンプする、Typescript以外のハイライトなどはいずれも未対応です。

仕組み

以下のような流れで実行しています。

  1. コードベースを構文解析してチャンクに分割する
  2. 分割したチャンクについてembeddingを求め、DBに格納する
  3. クエリをembeddingに変換し、ベクトル検索を行うことで類似したチャンクを発見する
  4. 以上を実装したコマンドをNeovimから呼び出す

開発の流れ

チャンキング

コードをチャンキングする方法はいくつかありますが、今回は自作のcode-chopperを使ってコードを分割しました。
https://github.com/sirasagi62/code-chopper

埋め込みモデルの選定

分割したチャンクについてテキスト埋め込みを生成する必要があります。テキスト埋め込みモデルはOpenAIやGoogleがあるので、それを使ってもよいですが、実はローカルで動かせるモデルもいくつか存在しています。今回は、コスト削減のためにローカルで埋め込みを生成しようと思います。

ローカルで動かせる代表的なモデルとしては

  • multilingual-e5(me5)
  • bge-m3
  • EmbeddingGemma
  • Qwen-Embedding-0.6B

あたりがあります。ただ、今回は通常のベクトル検索とは異なり、コード検索を日本語で、しかもローカルで行うため、

  1. 多言語(日本語含む)に対応していること
  2. コードでもembeddingが取れる
  3. ローカルで動かすので、パラメータ数が200M未満

を満たすのが望ましいです。

ところで、ローカルでコードをベクトル検索する事例は意外と少ないです。これは、近年のローカル埋め込みモデルが

  • SoTAを獲りに行くためにパラメータ数が1Bを越えるような場合が多い
  • コード検索がタスクに考慮されない場合が多い
  • 英語のみにしか対応していない場合が多い

といった傾向を持ち、通常のRAGよりも実装難易度が高いことが一因であると考えられます（そしてそもそも需要が少ない）。このため、OpenAIやVoyageのAPI経由でコードをベクトル検索する例は散見されますが、ローカルモデルについてはどのモデルが優れているか、といった議論がほとんどされていません。今回これを明らかにするために以下の実験を行いました。

コード埋め込みと日本語性能にある程度、優れているとされており、パラメータが200M以下の

の４つのモデル（いずれも8bitで量子化されたGGUF）について、llama.cpp付属のllama-serverを用いて試験しました。試験には以下のコードを用いました。

試験用コード
import json
import requests
import numpy as np
from scipy.spatial.distance import cdist

# -------------------------------------------------
# 1. テスト用コードと説明文
# -------------------------------------------------
code_snippets = [
    """def fibonacci(n):
    a, b = 0, 1
    for _ in range(n):
        a, b = b, a + b
    return a""",

    """import numpy as np
def normalize(v):
    norm = np.linalg.norm(v)
    return v / norm if norm != 0 else v""",

    """def is_prime(num):
    if num < 2:
        return False
    for i in range(2, int(num**0.5) + 1):
        if num % i == 0:
            return False
    return True"""
]
descriptions = [
    "フィボナッチ数列を計算する関数",
    "ベクトルを正規化するユーティリティ",
    "整数が素数かどうか判定する関数"
]
# -------------------------------------------------
# 2. llama.cpp の埋め込みエンドポイントへリクエスト
# -------------------------------------------------
EMBED_URL = "http://127.0.0.1:8080/embedding"

def get_embedding(text: str) -> np.ndarray:
    """テキストをサーバーに送ってベクトルを取得する"""
    payload = {"content": text}
    headers = {"Content-Type": "application/json"}
    resp = requests.post(EMBED_URL, data=json.dumps(payload), headers=headers)
    resp.raise_for_status()
    return np.array(resp.json()[0]["embedding"], dtype=np.float32)

# -------------------------------------------------
# 3. 全テキストの埋め込み取得
# -------------------------------------------------
code_embeds = [get_embedding("passage: "+c) for c in code_snippets]
desc_embeds = [get_embedding("query: "+d) for d in descriptions]

def most_similar(desc_vec: np.ndarray, code_vecs: list[np.ndarray]) -> int:
    desc_2d = desc_vec.reshape(1, -1)

    code_2d = np.vstack(code_vecs)

    distances = cdist(desc_2d, code_2d, metric="cosine")[0]
    print("dist:",distances)
    return int(np.argmin(distances))

# -------------------------------------------------
# 4. 結果表示
# -------------------------------------------------
for i, d in enumerate(descriptions):
    idx = most_similar(desc_embeds[i], code_embeds)
    print(f"説明: {d}")
    print("最も近いコード:")
    print(code_snippets[idx])
    print("-" * 40)

試験の結果、3つのembeddingについて正しい説明文を選択できたのが、graniteのみであったため、今回はgraniteを採用しました。

なお、多言語性能についてはgranite = me5 >= jina > nomic,コード性能についてはgranite >> jina > me5 = nomicという印象です。

埋め込みモデルの呼び出し

code-chopperがJSライブラリであるため、モデルの呼び出しにはtransformers.jsを利用しています。transformers.jsはONNXモデルの呼び出しにのみ対応しており、graniteは公式でONNXを提供しています。しかし、

  • transformers.jsのONNXモデルは周辺ファイルの配置含めて、やり方が決まっておりgranite公式はそれに則っていない
  • 量子化モデルが提供されない

という理由があり、公式ONNXを利用するのは困難です。というわけで、独自にONNXに変換し、これを呼び出しています。
https://huggingface.co/sirasagi62/granite-embedding-107m-multilingual-ONNX

検索機能の実装

ベクトル検索にはsqlite-vecを利用しました。sqlite-vecはsqliteが使える環境ならどこでも動く上、出力が単独ファイルで済むので、RAGを実装する上で非常に便利です。
https://github.com/asg017/sqlite-vec

sqliteのアダプタとしては、bunの内蔵sqliteライブラリを使いました。Nodeのbetter-sqlite3と一定の互換性があるため、Nodeでもbetter-sqlite3を利用することでほとんど改造なしで動かせると思います。

UIの実装

実は一番苦労したのがNeovimでのUIの実装です。当初はtelescopeの拡張機能という形で実装するつもりだったのですが、telescopeが一覧取得→検索→選択、という流れで動作するためベクトル検索との相性が悪く、独自UIでの実装に舵を切りました。

UIライブラリとしてnui-components.nvimを利用しました。リアクティブに実装でき、かなり使いやすい印象でしたが、promptコンポーネントがドキュメントの通りに動作しなかったり、そもそも利用例が少なかったりと、これはこれで意外と大変でした。もっとも、これ以外のライブラリを使った場合、最低限の動作まで持っていくことが難しかったと思います。

chopgrepで言語関連の情報を扱うのを忘れていたため（code-chopper側で情報自体は取得できる）、今回はPreview画面の言語設定をtypescriptで決め打ちしています。

まとめ

本記事ではコードを自然言語で検索するコマンドchopgrepとそれを用いたchopgrep.nvimについてまとめました。特に本記事で扱ったコードを日本語でベクトル検索する事例については、比較記事や実装記事が多くないので、ぜひ参考にしていただければと思います。

Discussion