RAGの機能もGPT-4 Turboに丸投げすれば「ソースコード自動QA対応システム」ができるかもしれない

# はじめに

LLMの回答精度を向上させるためのアプローチとして主流となっているものに「RAG：Retrieval-Augmented Generation」があります。
外部の知識ベースから、重要な部分を抽出してプロンプトに入れることでユーザの質問に的確に答えられる仕組みですが、類似検索のアルゴリズムによって直接は関係のない知識を持ってきてしまうなど、いくつか課題があることが知られています。目まぐるしく新しいアプローチが提案されているものの、手順が複雑ですぐに取り入れるのが難しい場合もあるのが現状です。

当法人の以下の記事もぜひご覧ください！↓

https://qiita.com/DeepMata/items/3c27394e4475b1b7e7ff

https://qiita.com/DeepMata/items/d54bfdcb732663dfd01b

そこで、**今回はRAGに相当する機能もGPT-4 Turboに任せてしまおう**、というモチベーションで、主にソースコードの分析を行うシステムを開発してみました。

GPT-4 Turboであれば、もはやソースコード全部を一度に投げることもできてしまうかもしれませんが、プロンプトの長さにある程度の制限があるLLMにおいても応用することを考慮したアプローチにしています。


# やりたいこと

- **RAGの機能もGPT-4 Turboに任せることで、より精度の高い情報を抽出する**
- 複数のスクリプトファイル（Python、Java、C++など）で構成されたソフトウェア(GitHubなどに公開されているもの等)について、**機能追加やプログラム修正などユーザの要望に沿ってアドバイスしてくれるシステム**を作りたい

:::note warn
GitHubなどで公開されているソフトウェアについて利用することを前提としています。機密情報を含むソースコードでは行わないでください。
:::

# 全体像

仕組みは単純で、以下の3ステップから構成されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1373557/eec2d745-3664-3049-3067-5adb38b1b0d3.png)

① ソースコードをファイル単位で日本語に要約し、要約ナレッジとして蓄積
② ユーザの質問と要約ナレッジから、その質問の回答のために重要と思われるソースコードのファイルをリストアップ
③ ②で得られたファイルのソースコードの中身をプロンプトに入れ、具体的な操作アドバイスを出力させる

# 実装

実装全体は以下のリポジトリで公開しています。

https://github.com/keisuke-okb/code-qa-gpt4


## ① ソースコードをファイル単位で日本語に要約し、要約ナレッジとして蓄積

各ファイルごとに処理内容を最大200文字程度で要約させ、日本語で処理内容を簡潔に説明してもらいます。

```python
import os
import argparse
import glob
from tqdm import tqdm
from openai import OpenAI

files = sorted([f for f in glob.glob(os.path.join(args.input_dir, '**/*'), recursive=True) 
            if os.path.splitext(os.path.basename(f))[1] in target_ext_dict.keys()
            ])

if not args.recursive:
    files = [f for f in files
            if not "\\" in str(f).replace(os.path.join(args.input_dir, ""), "")
            ]

summary_path = os.path.join(args.input_dir, "summary.txt")

if not os.path.exists(summary_path):
    print(f"Target files count: {len(files)}")
    print(files)
    input("Press enter key to continue.")
    
    abstracts = ""
    for file in tqdm(files):
        with open(file, "r", encoding="utf-8") as f:
            texts = "".join(f.readlines())

        if len(texts.replace(" ", "")) > args.min_length:

            abs_length = int(len(texts.replace(" ", "")) // 20 // 100 * 100)
            rel_path = str(file).replace(os.path.join(args.input_dir, ""), "")

            # GPT-4で要約する
            ext = target_ext_dict[os.path.splitext(os.path.basename(file))[1]]
            prompt = f"""以下は、{ext}のスクリプトです。どのような処理を行っているのか最大でも200字程度で簡潔に要約してください。
このプログラムだけでは具体的な処理が不明な場合は推測してください。
また、複数のクラスや関数が定義されている場合は、それぞれの処理内容を簡単にリストアップしてください。
- ```{ext}
{texts}
- ```
"""
            response = client.chat.completions.create(
            model="gpt-4-1106-preview",
            messages=[
                    {"role": "user", "content": prompt}, 
                ],
                max_tokens=args.max_token,
            )
            summary = response.choices[0].message.content
            abstracts += f"【{rel_path}】\n{summary}\n\n\n"

    with open(summary_path, "w", encoding="UTF-8") as f:
        f.writelines(abstracts)

```

## ② ユーザの質問と要約ナレッジから、その質問の回答のために重要と思われるソースコードのファイルをリストアップ

```python
prompt = f"""以下は、スクリプトのファイル一覧と処理の説明です。
下の質問に答えるために最も関連すると思われるファイルを、関連する順番に箇条書きでリストアップしてください。
理由や説明は必要なく、ファイル名の箇条書きのみ出力してください。

出力の例：
- hoge.py
- fuga/aaa/bbb.py
- piyo/ccc.py

----------

{abstracts}

----------
質問：{args.question}
"""
response = client.chat.completions.create(
    model="gpt-4-1106-preview",
    messages=[
            {"role": "user", "content": prompt}, 
        ],
        max_tokens=args.max_token,
    )
res = response.choices[0].message.content
```


## ③ 質問に関連するソースコードの中身をプロンプトに入れ、具体的な操作アドバイスを出力させる

②で出力されたファイル名の一覧から、そのファイルのコードを読み込んでプロンプトに入れ込みます。

```python
contents = ""
count = 0
for l in res.split("\n"):
    if "- " in l:
        file = l.strip().replace("- ", "")
        try:
            with open(os.path.join(args.input_dir, file), "r", encoding="utf-8") as f:
                texts = "".join(f.readlines())
            
            contents += f"【{file}】\n{texts}\n\n\n"
            count += 1
            if count == args.max_rag_files:
                break
        except:
            print(f"Failed to load: {file}")

prompt = f"""あなたは、優秀なプログラマーアシスタントです。

以下は、プログラムのファイル名とその内容を並べたものです。
これらのプログラムをもとに、下の質問になるべく簡潔に答えてください。
----------

{contents}

----------
質問：{args.question}
"""
response = client.chat.completions.create(
    model="gpt-4-1106-preview",
    messages=[
            {"role": "user", "content": prompt}, 
        ],
        max_tokens=args.max_token_qa,
    )
res = response.choices[0].message.content
print(res)

```

# 動作例

今回は「StyleGAN2-PyTorch」のソースコードを対象に、いくつか質問を投げてみました。
※今回は簡易化のため、サブディレクトリは無視しています。

https://github.com/rosinality/stylegan2-pytorch

## 得られた要約ナレッジの抜粋

【apply_factor.py】
このPythonスクリプトは、指定されたStyleGAN2モデル（ジェネレーター）を使用して、潜在空間に沿った画像のサンプルを生成しています。コマンドライン引数としてモデルのパラメーターを受け取り、与えられた固有ベクトルの方向に沿って潜在ベクトルを動かし、変換された画像を生成し、保存します。

- argparseを使用してコマンドライン引数を解析
- ジェネレーターモデルをロードし、指定されたチェックポイントから重みをセット
- 潜在空間におけるトランケーションの平均を計算
- 指定された数のランダム潜在ベクトルを生成
- 固有ベクトルを使用して、元の潜在ベクトルを変形させることで画像のバリエーションを作成
- 変形された潜在ベクトルを用いて、ジェネレーターモデルで画像を生成
- 元の画像と変形された画像を並べて一つの画像にし、指定のプレフィックスとパラメーターを元にしたファイル名で保存

この処理は、StyleGAN2モデルが生成する潜在空間の特定の方向（因子）に沿って画像がどのように変化するかを可視化するためのものであり、機械学習またはGANの研究において使われることが想定されます。


【calc_inception.py】
このPythonスクリプトは、データセットからイメージサンプルの特徴量を抽出して、それらの統計量（平均と共分散）を計算し、結果をPickleファイルに保存する処理を行っています。利用されるネットワークはInception v3モデルです。argparseを使ってコマンドライン引数を処理し、データセット、イメージサイズ、バッチサイズ、サンプル数、ランダムフリップの有無を指定できます。特徴量はGPU上で抽出され（CUDAが利用可能な場合）、結果はCPUに転送されます。最終的な統計量は指定されたサンプル数に基づいて計算されます。

以下、定義されているクラスや関数のリストです：

クラス：
- `Inception3Feature`: Inception3モデルを継承し、イメージをInception v3モデルに適したサイズにリサイズして前処理し、特徴量を抽出する。

関数：
- `load_patched_inception_v3()`: パッチド（変更された）Inception v3モデルをロードする。
- `extract_features(loader, inception, device)`: DataLoaderからイメージを読み込み、Inceptionモデルを使って特徴量を抽出し、それらのリストを返す。

スクリプトの実行部分（`if __name__ == "__main__":`以下）では、適切なコマンドライン引数を解析し、処理を実行し、結果をPickleファイルに保存しています。


## Q1：潜在変数をファイルをして書き出すにはプログラムをどう変更すればよいですか？ 

GPT-4 Turboの回答：

`projector.py` ファイルには、画像を StyleGAN2 の潜在空間に投影し、その潜在変数をファイルに保存する機能が含まれています。具体的には、次の行で潜在変数やノイズなどがファイルに保存されています。

```python
torch.save(result_file, filename)
```

この `result_file` 辞書には、各入力画像に対応する潜在変数 (`"latent"`) とノイズ (`"noise"`) が含まれています。

特定の潜在変数をファイルに保存したい場合は、以下の手順に従います。

1. `result_file` 辞書にはすでに潜在変数が含まれていますので、特定のキーの潜在変数を取り出します。
2. 取り出した潜在変数を `.cpu().numpy()` メソッドを使って NumPy 配列に変換します。
3. NumPy ライブラリを使用して、潜在変数を `.npy` ファイルとして保存します。

具体的なコードは以下の通りです。

```python
# result_file 辞書から特定の入力名に対応する潜在変数を取り出す
latent_var = result_file[input_name]["latent"]

# Device (CUDA) から CPU に移動し、NumPy 配列に変換する
latent_numpy = latent_var.cpu().numpy()

# NumPy 配列をファイルに保存する
np.save(f"{input_name}_latent.npy", latent_numpy)
```

このコードを `projector.py` の `for i, input_name in enumerate(args.files):` ループの中に追加します。そうすることで、各入力画像に対応する潜在変数が NumPy ファイルとして個別に保存されます。

完全なコードの例を以下に示します：

```python
import numpy as np  # 必要に応じてインポート

# 既存のループ内に次のコードを追加
for i, input_name in enumerate(args.files):
    # ... 既存のコード ...

    # 潜在変数を NumPy 配列に変換して保存
    latent_var = result_file[input_name]["latent"]
    latent_numpy = latent_var.cpu().numpy()
    np.save(f"{input_name}_latent.npy", latent_numpy)
```

このコードを追加することで、各入力画像の潜在変数を NumPy ファイルとして保存できるようになります。


## Q2: 生成された画像の保存パスを好きな場所に変えたいとき、どのプログラムを変えればよいですか？ 

GPT-4 Turboの回答：

画像の保存パスを変更したい場合、変更すべきプログラムは以下の2つです。

1. 【generate.py】プログラムでは、`utils.save_image`関数の呼び出し部分で保存パスが`f"sample/{str(i).zfill(6)}.png"`と指定されています。このパスを好きな場所に変更したい場合は、この文字列を変更します。

2. 【apply_factor.py】プログラムでは、`utils.save_image`関数でファイル名のプレフィックスとして`f"{args.out_prefix}_index-{args.index}_degree-{args.degree}.png"`を使用しています。保存先のディレクトリも含めて好きなファイルパスに変更することができます。

上記のプログラムのどちらか、もしくは両方の`utils.save_image`関数の呼び出し部分で指定されているパスを変更することで、画像の保存パスを変更することができます。



# おわりに

今回は単純なアルゴリズムでしたが、ナレッジからの情報抽出の部分もGPT-4 Turboに任せることでソースコードのQA対応を行うシステムを作ってみました。
質問をいくつか投げてみましたが、「どのファイルを」「どう変更すればよいか」という観点で具体的にアドバイスしてもらえることが分かりました。


この動作例では要約ナレッジが長すぎて全体では約1万文字まで膨れてしまったため、要約の長さを短くすればより大規模なソフトウェアに対しても対応可能であると思います。

最後までお読みいただきありがとうございました！