RBS meets LLMs - 大規模言語モデルを用いた型推論を試してみた
Leaner 開発チームの黒曜(@kokuyouwind)です。
RubyKaigi 2023 で「RBS meets LLMs - Type inference using LLM」と題して LT させていただきました。
そのうち本編のアーカイブ映像が出るはずですが、文字でも読めるよう大まかな内容を記事にまとめておきます。
また時間に収まりきらなかった内容も少し触れておくので、本編見ていただいた方も 本編に盛り込めなかった話 の節を呼んでもらえると嬉しいです。
発表スライド
以下の slides.com のものがマスターです。
フォントが中華フォントに化けてしまっていますが、一応 SpeakerDeck に上げたものも埋め込んでおきます。[1]
発表概要
Ruby の型定義ファイルである RBS を、 ChatGPT で推測できるか試してみました。
特に変数名や関数名の単語の意味を使うことで、例えば name
という変数であれば String
の可能性が高そうなど、人間が行うような推測を行えるのではないか、という予測に基づいていくつか実験しています。
UserBuilder の型の推測
試しに、以下のような Ruby コードを考えてみます。
user = UserBuilder.new.name('kokuyouwind').build
本来はこれだけの情報量だと UserBuilder
の型の推測はほとんどできないですが、 Builder
のデザインパターンの知識を持っていれば「name
では UserBuilder
自身を返し、 build
で User
を返すはず」と推測できます。
これを以下のプロンプトで ChatGPT に問い合わせました。なおモデルは chatgpt-3.5-turbo
で Temperature: 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
を返して name
は Company
のメソッドではないかと想像できます。
ここでポイントとなるのが、 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#find
で Company
を返し、 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.rbuser = UserBuilder.new.name('kokuyouwind').build
assistant:
user_builder.rbsclass UserBuilder def name: (name: String) -> UserBuilder def build: () -> User end
user.rbs
class User attr_reader name: String end
user:
test.rbcompany = CompanyRepository.new.find(1).name
これに対する ChatGPT の回答は以下のようになります。
assistant:
company_repository.rbsclass 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 theTester
class would be:class Tester def hello: (name: untyped) -> String end
This is because the
hello
method takes in an argumentname
of any type and returns a string with the value ofname
interpolated.
name
が String
ではなく untyped
なのがちょっと微妙ですが、少なくとも hello
メソッドの存在が推測されているのは注目に値します。
これができるのであれば Rails などのライブラリで多用される動的メソッド定義も推測できる可能性が高いといえます。もし学習時の知識にないものでも、 FewShot やファインチューニングで対応できるのではないでしょうか。
既存の RBS の改善
今回はインパクト重視とスライド内の文量節約のため短いコードからの型推論を試しましたが、実用上は rbs prototype
などで下敷きとなる RBS を生成したうえで、 untyped
になっている箇所を具体的な型に置き換えていく作業をするケースが多いでしょう。
例えば、以下の 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
が得られます。
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.rbclass 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 aString
argument and returnsUserFactory
object itself, so that it can be chained with other methods.build
method returns aUser
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.cpp や ChatRWKV など、省メモリで動作する LLM の選択肢がいくつかあります。[3]
さらにオープンソース LLM の研究が盛んで今後も発展が予想されるため、機を見てローカル完結でどこまでの推論が行えるかも試していきたいです。
まとめ
いくつかの例で見ていただいたとおり、小さい例ではわずかな情報量でもかなり正確に型を推測してくれることがわかりました。
RubyKaigi 2023 の Matz Keynote でも触れられていましたが、 Ruby の型付けを自動で行う未来が決して夢物語ではなくなってきたと感じます。
もちろん既存の他アプローチが必要なくなるということはなく、むしろ特徴が異なるこれらを積極的に併用していけるようなツール設計が必要になるのではないでしょうか。
我ながら面白い結果が得られたので、引き続きツール化やプロンプトの改善に取り組んでいく予定です。
-
PDF ダウンロード時にフォントが埋め込まれていないっぽいです。こればかりはプラットフォームの仕様なのでしょうがない… ↩︎
-
何人かから「なんで ChatGPT-4 を使わないの?」と質問を受けました。端的には Waitlist 登録から進んでおらず API が使えなかったからですが、将来的にローカルの LLM が主流になりえることを考えると、賢すぎる ChatGPT-4 を使うより ChatGPT-3.5 くらいの賢さで実現可能な範囲を探るほうが実用的かなと思い、今回はこちらを使っています。 ↩︎
-
モデルと動作環境や、 transformer と RNN を区別せず議論してますが、ローカルで動かせる省メモリ LLMs 技術の例としてひっくるめています、ご了承ください。 ↩︎
Discussion