BigQueryでLLM（RAGを用いたテキスト生成）を使う

Google Cloud Next Tokyo2023でBigQueryからVertexAIのLLMを使用する話を聞いてきたので試してみる。
質問に対して、グラウンディング用データを参照しつつLLMで回答するということをSQLで実装できた。

全体的にはこんなイメージ。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/97c0decf-8f2e-4ad3-0902-5628ba92529c.png)

# 準備
## BigQuery Connection API を有効にする
https://cloud.google.com/bigquery/docs/working-with-connections?hl=ja

## 外部接続の作成
BigQueryの画面の左上のエクスプローラの追加をクリック
![Screenshot (94).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/59322142-3c67-4b01-5b6c-fdadb073476c.png)

外部データソースへの接続をクリック
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/9d953e1c-ebdb-e785-5844-ae59259ddbb2.png)

以下を設定して接続を作成
 - 接続タイプ：BigLakeとリモート関数
 - 接続ID：llm-conn
 - ロケーション：マルチリージョンでUS（対象のデータセットと同じロケーション？）

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/b9ed9592-64d6-9065-7091-30a3faa844fa.png)


外部接続が作成された
![Screenshot (95).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/6980d2b6-0a42-9ba2-797f-e273b992d73e.png)

## 権限設定
外部接続で作成されたサービスアカウント（外部接続画像のサービスアカウントID）に[Vertex AI ユーザー] ロール権限を付与する
![Screenshot (96).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/40fae47a-e60f-5537-3bae-30c9c16ca37f.png)

# モデルを作成する
以下を使ってみる
LLM：リモートサービスのCLOUD_AI_LARGE_LANGUAGE_MODEL_V1
Embedding：エンドポイントのtextembedding-gecko-multilingual

```sql:LLMのモデル作成
CREATE OR REPLACE MODEL `project_id.dataset.model_name`
  REMOTE WITH CONNECTION `us.llm-conn`
  OPTIONS (remote_service_type = 'CLOUD_AI_LARGE_LANGUAGE_MODEL_V1');
```

```sql:Embedのモデル作成
CREATE OR REPLACE MODEL `project_id.dataset.model_name_embedding`
  REMOTE WITH CONNECTION `us.llm-conn`
  OPTIONS (ENDPOINT = 'textembedding-gecko-multilingual');
```

# グラウンディングを用いた情報検索をする
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/fd28b673-06ec-37e2-b894-f1ba26219356.png)


要は自社のデータ等をLLMに渡して、LLMの学習していないデータを用いて回答してもらう方法

## グラウンディング用のデータ（埋め込み、Embedding）を作成する
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/4419460a-3b4b-84ce-5433-a200ce638287.png)


要は質問に似通ったデータを抽出しやすくするためのプロセスで、テキストデータを数値化（ベクトル化）する

まず以下のような色々な文章が記載されている項目（text列）をもつテーブルを作成する
（今回はChatGPTに小売店でのサンプルFAQを作ってもらった）
![Screenshot (97).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/bca9aecc-c09c-e482-af66-50036e93dbd6.png)


埋め込み列をもつテーブルを作成する
埋め込みを作成したい列はcontentというエイリアスにしておく必要がある
元々contentという列の名前をもつテーブルであればテーブルの指定だけでも良い

```sql:埋め込み列をもつテーブル作成
create table `project_id.dataset.table_name_embedding`
as
(
select *
from
  ML.GENERATE_TEXT_EMBEDDING(
    MODEL `project_id.dataset.model_name_embedding`,
    (
     select 
       text as content
     from 
       `project_id.dataset.table_name_embedding` 
    ),
    STRUCT(TRUE AS flatten_json_output) --jsonデータを展開して列にするかのオプション
  )
);
```

こんな感じのtext_embedding列を持つテーブルが作成された
![Screenshot (98).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/1e601920-9aa2-311b-0cd8-a4eb3a3c96bd.png)


## 埋め込んだテキストのk-meansモデルを作成する
埋め込みを 10 個のクラスタに分類するk-meansモデルを作成する
いきなり埋め込み同士の距離計算するより、埋め込みの近いクラスタのデータのみで距離計算した方がパフォーマンスが良いらしい

```sql:埋め込みのk-meanモデル作成
CREATE OR REPLACE MODEL `project_id.dataset.model_name_cluster`
OPTIONS (
  model_type = 'KMEANS',
  KMEANS_INIT_METHOD = 'KMEANS++',
  num_clusters = 10) AS (
    SELECT
      text_embedding
    FROM
      `project_id.dataset.table_name_embedding`
);
```

## 質問に近いデータを抽出する
### 質問内容を埋め込みし、質問内容が属するクラスタを予測する
```sql:
WITH query_test as
(
SELECT
  *
FROM
  ML.GENERATE_TEXT_EMBEDDING(
    MODEL `project_id.dataset.model_name_embedding`,
    (SELECT
      "現金以外の支払いはできますか" AS content),
      STRUCT(TRUE AS flatten_json_output))
)
SELECT
  centroid_id,
  text_embedding
FROM
  ML.PREDICT(
    MODEL `project_id.dataset.model_name_cluster`,
    (SELECT
       text_embedding
     FROM
       query_test)
  );
```

こんな感じにクラスタIDが7で埋め込まれた結果が返ってくる
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/36a8e10e-74c8-88cb-d6a2-700342625486.png)


### グラウンディング用のデータから同一クラスタのデータを抽出する

```sql:
SELECT
    *
FROM
  ML.PREDICT(
    MODEL `project_id.dataset.model_name_cluster`,
    (SELECT
       text_embedding,
       content
     FROM
       `project_id.dataset.table_name_embedding`)
  )
  WHERE centroid_id = 7
  ;
  ```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/a66b4c09-6f0f-0151-5527-806b16fad31e.png)


### クラスタ内で質問内容との距離を計算する
ここまでの内容を組み合わせて、同一クラスタ内で、質問と最初に作成したテーブルの文章で埋め込みのベクトルの距離（コサイン類似度）を計算する

```sql
WITH 
--質問文の埋め込み
query_test as
(
SELECT
  *
FROM
  ML.GENERATE_TEXT_EMBEDDING(
    MODEL `project_id.dataset.model_name_embedding`,
    (SELECT
      "現金以外の支払いはできますか" AS content),
      STRUCT(TRUE AS flatten_json_output))
),
--質問文のクラスタ抽出
query_cluster as
(
SELECT
  centroid_id,
  content,
  text_embedding
FROM
  ML.PREDICT(
    MODEL `project_id.dataset.model_name_cluster`,
    (SELECT
       text_embedding,
       content
     FROM
       query_test)
  )
),
--グラウンディング用データから質問文と同じクラスタを抽出
answer_cluster as
(
SELECT
    *
FROM
  ML.PREDICT(
    MODEL `project_id.dataset.model_name_cluster`,
    (SELECT
       text_embedding,
       content
     FROM
       `project_id.dataset.table_name_embedding`)
  )
  WHERE centroid_id in (select centroid_id from query_cluster)
)
--質問文と同一のクラスタのグラウンディングデータから質問文と距離が近い順に出力する
SELECT
  s.content AS search_content,
  c.content AS content,
  ML.DISTANCE(s.text_embedding, c.text_embedding, 'COSINE') AS distance
FROM
  query_cluster AS s,
  answer_cluster AS c
ORDER BY
  distance ASC
;
```

それっぽい結果が返ってきた

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/866bf539-631d-33e9-6c44-7acf0a52359a.png)


## LLMに回答させる
質問内容とグラウンディング用の回答からLLMに回答させてみる

```sql:
SELECT 
  prompt as query,
  string(ml_generate_text_result.predictions[0].content) as answer
FROM
  ML.GENERATE_TEXT(
    MODEL `project_id.dataset.model_name`,
    (
      SELECT
        CONCAT(
          '質問内容に以下の情報を用いながら回答してください。',
          '質問内容：','現金以外の支払いはできますか',
          ' 情報：', '店舗でのクレジットカード支払いは可能ですか？はい、主要なクレジットカードでの支払いを受け付けています。') AS prompt
    ),
    STRUCT(
      0.1 AS temperature,
      1000 AS max_output_tokens,
      0.1 AS top_p,
      10 AS top_k));
```

それっぽく返ってきた
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/bd6927a0-8132-2af2-9f9c-24c59ad0c3db.png)



# 一連の流れをプロシージャにする
ここまでの流れをプロシージャにまとめて、質問内容を引数に渡してLLMからの回答をもらえるようにする

pythonから使用することも考慮し、OUT引数も設定しておく

```sql:
CREATE OR REPLACE PROCEDURE `project_id.dataset.procedure_name`(
  query_string STRING,
  OUT llm_answer STRING
)
BEGIN
  DECLARE grounding STRING;
  DECLARE prompt_pre STRING DEFAULT '質問内容に以下の情報を用いて自然な文章で回答してください。' ;
  DECLARE prompt_suf STRING DEFAULT 'なるべく詳細に回答してください。';
  
  -- 質問からグラウンディングデータに一番近いものを取得する
  SET grounding = (
  WITH query_test as
      (
        SELECT
          *
        FROM
          ML.GENERATE_TEXT_EMBEDDING(
            MODEL `project_id.dataset.model_name_embedding`,
            (SELECT
              query_string AS content),
              STRUCT(TRUE AS flatten_json_output))
      ),
      query_cluster as
      (
        SELECT
          centroid_id,
          content,
          text_embedding
        FROM
          ML.PREDICT(
            MODEL `project_id.dataset.model_name_cluster`,
            (SELECT
               text_embedding,
               content
             FROM
               query_test)
          )
        ),
      answer_cluster as
      (
        SELECT
          *
        FROM
          ML.PREDICT(
            MODEL `project_id.dataset.model_name_cluster`,
            (SELECT
               text_embedding,
               content
             FROM
               `project_id.dataset.table_name_embedding`)
          )
        WHERE centroid_id in (select centroid_id from query_cluster)
      )
      --質問とグランディング用データで一番距離の近いデータを抽出する
      SELECT
        c.content AS content,
      FROM
        query_cluster AS s,
        answer_cluster AS c
      QUALIFY
        ROW_NUMBER()OVER(order by ML.DISTANCE(s.text_embedding, c.text_embedding, 'COSINE') ASC) = 1) ;
  
  
  --LLM回答用のデータを取得する
  SET llm_answer = (
      SELECT 
        string(ml_generate_text_result.predictions[0].content) as answer
      FROM
        ML.GENERATE_TEXT(
          MODEL `project_id.dataset.model_name`,
          (
            SELECT
              CONCAT(
                prompt_pre,
                '質問内容：',query_string,
                ' 情報：', grounding,
                prompt_suf) AS prompt
          ),
          STRUCT(
            0.1 AS temperature,
            1000 AS max_output_tokens,
            0.1 AS top_p,
            10 AS top_k))
  );

  
  SELECT llm_answer;
  
END;
```

プロシージャを実行する
```sql
DECLARE result STRING;
call `project_id.dataset.procedure_name`('返品はできますか。', result);
```

それっぽく返ってきた
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/209773/9b398aac-dd91-46e1-9dc4-ce902437bd85.png)


## pythonでストアドの返り値を取得する
google colabでストアドの返り値を取得するには以下の感じで、ストアド実行後にOUT引数をSELECTする

```python
from google.colab import auth
auth.authenticate_user()
from google.cloud import bigquery
client = bigquery.Client(project='<project_id>')

# ストアドプロシージャを実行するクエリ
query = """
DECLARE result STRING;
call `project_id.dataset.procedure_name`('返品はできますか。', result);
SELECT result;
"""

# クエリの実行
query_job = client.query(query)

# 結果の取得
results = query_job.result()
for row in results:
    print(row[0])
```

# 参考
[BigQuery で Vertex AI モデルを使って、SQL のみで LLM のテキスト生成を実行する](https://cloud.google.com/blog/ja/products/ai-machine-learning/llm-with-vertex-ai-only-using-sql-queries-in-bigquery)
[The CREATE MODEL statement for remote models](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-create-remote-model)
[ML.GENERATE_TEXT_EMBEDDING 関数を使用してテキストを埋め込む](https://cloud.google.com/bigquery/docs/generate-text-embedding?hl=ja)
[テキスト エンベディングを使用して基本的なセマンティック検索を行う](https://cloud.google.com/bigquery/docs/text-embedding-semantic-search?hl=ja)
[The ML.DISTANCE function](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-distance?hl=ja)
[SQL ストアド プロシージャを操作する](https://cloud.google.com/bigquery/docs/procedures?hl=ja#pass_a_value_out_with_an_output_parameter)