💎

RBS meets LLMs - 大規模言語モデルを用いた型推論を試してみた

2023/05/19に公開

Leaner 開発チームの黒曜(@kokuyouwind)です。

RubyKaigi 2023 で「RBS meets LLMs - Type inference using LLM」と題して LT させていただきました。

https://rubykaigi.org/2023/presentations/lt/

そのうち本編のアーカイブ映像が出るはずですが、文字でも読めるよう大まかな内容を記事にまとめておきます。

また時間に収まりきらなかった内容も少し触れておくので、本編見ていただいた方も 本編に盛り込めなかった話 の節を呼んでもらえると嬉しいです。

発表スライド

以下の slides.com のものがマスターです。

https://slides.com/kokuyouwind/rbs-meets-llms

フォントが中華フォントに化けてしまっていますが、一応 SpeakerDeck に上げたものも埋め込んでおきます。[1]

発表概要

Ruby の型定義ファイルである RBS を、 ChatGPT で推測できるか試してみました。

特に変数名や関数名の単語の意味を使うことで、例えば name という変数であれば String の可能性が高そうなど、人間が行うような推測を行えるのではないか、という予測に基づいていくつか実験しています。

UserBuilder の型の推測

試しに、以下のような Ruby コードを考えてみます。

user = UserBuilder.new.name('kokuyouwind').build

本来はこれだけの情報量だと UserBuilder の型の推測はほとんどできないですが、 Builder のデザインパターンの知識を持っていれば「name では UserBuilder 自身を返し、 buildUser を返すはず」と推測できます。

これを以下のプロンプトで ChatGPT に問い合わせました。なおモデルは chatgpt-3.5-turboTemperature: 0 としています。[2]

system:
Act as Ruby type inferrer; only output RBS when Ruby code is input. Use class names, variable names, etc., to infer type.

user:

user = UserBuilder.new.name('kokuyouwind').build

これに対する ChatGPT の返答は以下です。

assistant:
Based on the code snippet, it seems that there is a UserBuilder class that has a name method which takes a string argument and returns an instance of UserBuilder. The UserBuilder class also has a build method which returns an instance of User.

Therefore, we can infer the following RBS type signature:

class User
  attr_reader name: String
end

class UserBuilder
  def name(name: String) -> UserBuilder
  end

  def build -> User
  end
end

RBS としての文法は全く正しくないのですが、 UserBuilder#name では UserBuilder を返し、 UserBuilder#build では User を返すという重要なところを推測できています。

このことから、 ChatGPT を型の推測に使えそうだと考えられます。

CompanyRepository の型の推測

次に、同じような例で以下の Ruby コードを考えてみます。

company = CompanyRepository.new.find(1).name

今度は CompanyRepository なので、 find の時点で Company を返して nameCompany のメソッドではないかと想像できます。

ここでポイントとなるのが、 UserBuilder の例と今回の例は構文上 variable = SomeClass.new.method1(arg1).method2 という全く同じ形をしていることです。

変数名やクラス名を単なるラベル以上に利用しないアプローチでは、この 2 つの例では情報量が全く変わらず、したがって全く同じ型の推論しかできないはずのものになっています。

先ほどと同様のプロンプトで尋ねると、 ChatGPT の返答は以下のようになります。

Based on the code snippet, it seems that CompanyRepository is a class that has a method find which takes an integer argument and returns an object with a name attribute. Therefore, we can infer the following RBS type signature:

class CompanyRepository
  def find: (Integer) -> Company
end

class Company
  attr_reader name: String
end

Note that we are assuming that CompanyRepository returns an object of type Company. If this is not the case, the RBS type signature would need to be adjusted accordingly.

今回は RBS らしい文法になっている上、ちゃんと CompanyRepository#findCompany を返し、 Company#name メソッドが String を返すことも推測できています。

このことから、構文上同じ形のプログラムに対して「変数名やクラス名を型の推測に利用できているらしい」ことが確かめられました。

FewShot を用いた出力の改善

ここまでで、 ChatGPT には単語の意味を元にした型の推測ができるとわかりました。

しかしながら、いくつかの課題も見つかりました。

  • UserBuilder の例では RBS の文法が壊れていた
  • 前後に無駄な文を出力していることがあった
    • RBS のみを出力するほうがツールなどから利用しやすい
  • RBS が 1 つのコードにまとまって出力されていた
    • 実用上はクラスごとにファイルを分割したい

これらの課題を改善するため、 FewShot を利用します。これは模範解答を例示した上でメインの問い合わせを渡すことで、模範解答に沿った形式の回答をしてもらう技法です。

今回は UserBuilder での型推論の模範解答を挙げた上で、 CompanyRepository の型を推論しもらいます。

入力したプロンプトは以下の通りです。ここで、 assistant: となっているところも入力になっており、要は会話履歴を捏造していることに注意してください。

user:
Act as Ruby type inferrer; only output RBS when Ruby code is input. Use class names, variable names, etc., to infer type. If you understand, say ok.

assistant:
ok.

user:
test.rb

user = UserBuilder.new.name('kokuyouwind').build

assistant:
user_builder.rbs

class UserBuilder
  def name: (name: String) -> UserBuilder
  def build: () -> User
end

user.rbs

class User
  attr_reader name: String
end

user:
test.rb

company = CompanyRepository.new.find(1).name

これに対する ChatGPT の回答は以下のようになります。

assistant:
company_repository.rbs

class CompanyRepository
  def find: (id: Integer) -> Company
end

company.rbs

class Company
  attr_reader name: String
end

今回は RBS の文法にきっちり則った上で、さらに company_repository.rbs と company.rbs とで RBS を分割して出力しており、 FewShot で出力形式を望む形に整えられました。

本編に盛り込めなかった話

5 分という制限があったため盛り込めなかったのですが、他の例もいくつか試して割と良い結果が得られていたため、合わせてまとめておきます。

define_method に対する型の推測

静的な型解析アプローチは define_method のような動的定義を苦手としていますが、 ChatGPT ではある程度推測できるようです。

以下が入力プロンプトです。

system:
Act as Ruby type inferrer; only output RBS when Ruby code is input. Use class names, variable names, etc., to infer type.

user:

class Tester
  define_method :hello, -> { |name| "hello, ${name}." }
end

これに対する ChatGPT の出力が以下です。

assistant:
Based on the code snippet, the inferred RBS type signature for the Tester class would be:

class Tester
  def hello: (name: untyped) -> String
end

This is because the hello method takes in an argument name of any type and returns a string with the value of name interpolated.

nameString ではなく untyped なのがちょっと微妙ですが、少なくとも hello メソッドの存在が推測されているのは注目に値します。

これができるのであれば Rails などのライブラリで多用される動的メソッド定義も推測できる可能性が高いといえます。もし学習時の知識にないものでも、 FewShot やファインチューニングで対応できるのではないでしょうか。

既存の RBS の改善

今回はインパクト重視とスライド内の文量節約のため短いコードからの型推論を試しましたが、実用上は rbs prototype などで下敷きとなる RBS を生成したうえで、 untyped になっている箇所を具体的な型に置き換えていく作業をするケースが多いでしょう。

例えば、以下の user_factory.rb があるとします。

user_factory.rb
class UserFactory
  def name(name)
    @name = name
  end

  def build
    User.new(name: name)
  end
end

これに対して rbs prototype rb user_factory.rb を実行すると以下の user_factory.rbs が得られます。

user_factory.rbs
class UserFactory
  def name: (untyped name) -> untyped

  def build: () -> untyped
end

通常、ここで untyped になっている箇所は手作業で具体型に直していくのですが、これを ChatGPT で推測できるとかなり便利そうです。

以下のプロンプトで問い合わせてみましょう。

system:
Act as Ruby type inferrer; refine RBS when Ruby code is input. Output only RBS codes. Use class names, variable names, etc., to infer type.

user:
user_factory.rb

class UserFactory
  def name(name)
    @name = name
  end

  def build
    User.new(name: name)
  end
end

user_factory.rbs

class UserFactory
  def name: (untyped name) -> untyped

  def build: () -> untyped
end

これに対する ChatGPT の回答は以下です。

assistant:
Refined user_factory.rbs:

class UserFactory
  def name: (String name) -> UserFactory

  def build: () -> User
end

Explanation:

  • name method takes a String argument and returns UserFactory object itself, so that it can be chained with other methods.
  • build method returns a User object.

もともと untyped だった箇所が、妥当そうな型に置き換わっています。なお無駄な文章が出ているのは前記のとおり FewShow で補正できるはずです。

個人的にはこのユースケースの需要が最も高いと考えているので、様々な例でプロンプトを変えながら試して実用に耐えるのか検証したいところです。

今後やりたいこと

コマンドラインツール化

上記のとおり割と良い感じに推論できているのですが、作業が間に合わず ChatGPT API を BetterChatGPT から叩いて試すところまでで終わってしまいました。

各種のファイルで試すためにはコマンドラインツール化したほうが便利で、例えば rbs-refine --ruby user_factory.rb --rbs user_factory.rbs のようなコマンドで改善した user_factory.rbs が出力される、といったツールがあるとプロンプト改善するにしても他の人に使ってみてもらうにしても良いはずです。

これは結構簡単にできるはずなので、ちゃっとやって Gem 公開したいです。

ファインチューニング

今回は素の ChatGPT に推論してもらいましたが、おそらく学習データの中に RBS 関連のコードはさほど多くなく、なにより 2020 年以降に入った文法には対応していないことが予想されます。

このため Ruby コードとそこから推測される RBS コードをペアにしたデータを用意しファインチューニングすることで、もっと性能の良い推論ができるのではないかと期待しています。

ただデータの用意が大変なことと、料金的に Ada を使うのが妥当なのですがこの用途でどこまでの性能が出るのか未知数なこともあり、優先度は低そうです。

エージェントスタイルの自律駆動

例えば user_factory.rbs の型の詳細化を行う場合、 user_factory.rb を参考にするだけでなく、 user.rbs があるならその情報も使ったほうが正確に推論できるはずです。

これをするには「どの情報を使うとより正確な判断ができるか」を考える必要があり、 LangChainの Chain や Agent を使う必要がありそうだと考えています。

さらにより発展させて、例えば「プロジェクト全体の型情報を順に推論して、 steep check した結果がエラーにならない範囲で詳細化して」といった目標を元に、 Auto-GPT のような自律駆動で「必要なファイルの参照」「特定の RBS の詳細化」「steep check の実行と結果を元にした修正」を行うことも考えています。

これはかなり夢の話ですが、もし可能になればプロジェクトへの RBS 型チェック導入をほぼ自動化できそうです。

ローカル LLM での動作

開発ツールとして常用することを考えたり、上記のエージェントスタイルの自律駆動などを考えると、従量料金のかかる ChatGPT API の代わりにローカルで動作する LLM を利用する選択肢があると便利そうです。

現在でも llama.cppChatRWKV など、省メモリで動作する LLM の選択肢がいくつかあります。[3]

さらにオープンソース LLM の研究が盛んで今後も発展が予想されるため、機を見てローカル完結でどこまでの推論が行えるかも試していきたいです。

まとめ

いくつかの例で見ていただいたとおり、小さい例ではわずかな情報量でもかなり正確に型を推測してくれることがわかりました。

RubyKaigi 2023 の Matz Keynote でも触れられていましたが、 Ruby の型付けを自動で行う未来が決して夢物語ではなくなってきたと感じます。

もちろん既存の他アプローチが必要なくなるということはなく、むしろ特徴が異なるこれらを積極的に併用していけるようなツール設計が必要になるのではないでしょうか。

我ながら面白い結果が得られたので、引き続きツール化やプロンプトの改善に取り組んでいく予定です。

脚注
  1. PDF ダウンロード時にフォントが埋め込まれていないっぽいです。こればかりはプラットフォームの仕様なのでしょうがない… ↩︎

  2. 何人かから「なんで ChatGPT-4 を使わないの?」と質問を受けました。端的には Waitlist 登録から進んでおらず API が使えなかったからですが、将来的にローカルの LLM が主流になりえることを考えると、賢すぎる ChatGPT-4 を使うより ChatGPT-3.5 くらいの賢さで実現可能な範囲を探るほうが実用的かなと思い、今回はこちらを使っています。 ↩︎

  3. モデルと動作環境や、 transformer と RNN を区別せず議論してますが、ローカルで動かせる省メモリ LLMs 技術の例としてひっくるめています、ご了承ください。 ↩︎

リーナーテックブログ

Discussion