はじめに
Retrieval-Augmented Generation (RAG)をbeautifulsoupによるWebスクレイピングとDense Retriver(文章の意味的な類似度を反映)により実装します。
指定されたウェブページからコンテンツを抽出し、クエリに基づいて最も関連性の高い情報を取得し、取得した情報をクエリにくっつけることでプロンプトを作成します。
全部張り付ければよくね?? という疑問
わざわざクエリに基づいた最も関連性の高い情報を、、、などとせず、適当に全部張り付けて渡せばいいと考える方もいると思います。それに関しては、以下の論文が参考になると思います。
プロンプトの長さと性能の関係
Lost in the Middle: How Language Models Use Long Contexts
この論文によると、RAGで与えた文章の中で求めている情報が真ん中にある場合、むしろ情報を与えていないときよりも性能が落ちる可能性があることを示しています。
0. RAGを実装したコード
chunk_rag.py
import requests
from bs4 import BeautifulSoup
import re
# 文章や画像のembeddingを扱うためのライブラリ、HuggingFaceのモデルを使うこともできる
from sentence_transformers import SentenceTransformer
def generate_prompt():
# https://huggingface.co/Alibaba-NLP/gte-large-en-v1.5
emb_model = SentenceTransformer("Alibaba-NLP/gte-large-en-v1.5", trust_remote_code=True)
# URLの設定
url = "https://beautiful-soup-4.readthedocs.io/en/latest/"
# ページの取得
response = requests.get(url)
soup = BeautifulSoup(response.text, 'html.parser')
# データを格納するリスト
data_list = []
# メインコンテンツ領域を特定
main_content = soup.find('div', class_='document')
# 段落とコードブロックを順番に処理
elements = main_content.find_all(['p', 'div', 'pre'])
current_paragraph = ""
for element in elements:
if element.name == 'p':
# 段落のテキストを取得
paragraph_text = element.get_text().strip()
if paragraph_text:
current_paragraph = paragraph_text
elif element.name == 'div' and 'highlight-default' in element.get('class', []):
# コードブロックのテキストを取得
code_block = element.find('pre')
if code_block and current_paragraph:
code_text = code_block.get_text().strip()
# データリストに追加
data_list.append({
"paragraph": current_paragraph,
"code": code_text
})
current_paragraph = "" # 次の段落のために初期化
prompt = "how to install beautifulsoup4"
query_embeddings = emb_model.encode([prompt], normalize_embeddings=True)
input_texts = [str(data) for data in data_list]
passage_embeddings = emb_model.encode(input_texts, normalize_embeddings=True)
scores = (query_embeddings @ passage_embeddings.T) * 100
top_k = 3
top_k_idx = scores[0].argsort()[::-1][:top_k]
retrieved_text = f"""
* {input_texts[top_k_idx[0]]}
* {input_texts[top_k_idx[1]]}
* {input_texts[top_k_idx[2]]}
"""
prompt = f"""{retrieved_text}
上記の文章に基づいて、質問に日本語で回答してください。
質問: {prompt}
回答:"""
return prompt
1. 今回必要なライブラリのインポート
import requests
from bs4 import BeautifulSoup
import re
from sentence_transformers import SentenceTransformer
-
requests
: ウェブページの取得に使用 -
BeautifulSoup
: HTMLの解析に使用 -
SentenceTransformer
: テキストの埋め込み(エンコード)に使用
2. generate_prompt
関数
この関数がchunk_rag.py
の主要な処理を行います。
2.1 埋め込みモデルの初期化
emb_model = SentenceTransformer("Alibaba-NLP/gte-large-en-v1.5", trust_remote_code=True)
Alibaba NLPの事前学習済みモデルを使用して、テキストを高次元ベクトルに変換します。
2.2 ウェブページの取得と解析
url = "https://beautiful-soup-4.readthedocs.io/en/latest/"
response = requests.get(url)
soup = BeautifulSoup(response.text, 'html.parser')
Beautiful Soupのドキュメントページを取得し、BeautifulSoupオブジェクトを作成します。
urlを変更することで、別のサイトでもWebスクレイピングが可能です。
2.3 コンテンツの抽出
data_list = []
main_content = soup.find('div', class_='document')
elements = main_content.find_all(['p', 'div', 'pre'])
current_paragraph = ""
for element in elements:
# 段落とコードブロックの抽出ロジック
...
ページの主要なコンテンツ領域から段落とコードブロックを抽出し、data_list
に格納します。
2.4 クエリと文書の埋め込み
prompt = "how to install beautifulsoup4"
query_embeddings = emb_model.encode([prompt], normalize_embeddings=True)
input_texts = [str(data) for data in data_list]
passage_embeddings = emb_model.encode(input_texts, normalize_embeddings=True)
クエリと抽出したテキストを埋め込みベクトルに変換します。
2.5 類似度の計算と上位k件の取得
scores = (query_embeddings @ passage_embeddings.T) * 100
top_k = 3
top_k_idx = scores[0].argsort()[::-1][:top_k]
クエリと各テキスト間のコサイン類似度を計算し、最も関連性の高い上位3件を選択します。
2.6 プロンプトの生成
retrieved_text = f"""
* {input_texts[top_k_idx[0]]}
* {input_texts[top_k_idx[1]]}
* {input_texts[top_k_idx[2]]}
"""
prompt = f"""{retrieved_text}
上記の文章に基づいて、質問に日本語で回答してください。
質問: {prompt}
回答:"""
選択されたテキストを使用してプロンプトを構築し、言語モデルへの入力として使用します。
まとめ
今回のRAGの流れは以下の通りです。
- 関連ウェブページからコンテンツを抽出
- クエリと抽出したテキストを埋め込みベクトルに変換
- コサイン類似度に基づいて最も関連性の高いテキストを選択
- 選択されたテキストを使用してプロンプトを生成
この方法により、言語モデルは質問に関連する具体的な情報を参照しながら回答を生成できるようになり、より正確で文脈に即した応答が可能となります。
補足: その他のコード
RAGの実装以外に、モデルのセットアップや実行に関するコードを張っておきます。
default_use.py
import torch
import time
import os
import psutil
from init import model, tokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
import cProfile
from chunk_rag import generate_prompt
#text作成
def generate_text(prompt):
messages = [ { "role": "user", "content": prompt} ]
#tokenizer.apply_chat_template() メソッドは、チャットメッセージをモデル固有の形式に変換
input_ids = tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
padding=True,
#truncationはmax_lengthの中に切ってくれる。
truncation=True,
max_length=512,
return_tensors="pt"
).to(model.device)
#terminators は、生成を終了させるトークンIDのリスト
terminators = [
#tokenizer.eos_token_id は、文の終わりを示す標準的なトークンID
tokenizer.eos_token_id,
#tokenizer.convert_tokens_to_ids("<|eot_id|>") は、
#カスタムの終了トークン "<|eot_id|>" をトークンIDに変換
tokenizer.convert_tokens_to_ids("<|eot_id|>")
]
#eos_token_id=terminators は、先ほど定義した終了トークンのリストを指定
#これらのトークンのいずれかが生成されると、生成プロセスが終了
#勾配計算はする必要がないためno_gradを使用
# 開始時間を記録
start_time = time.time()
with torch.no_grad():
result = model.generate(
input_ids,
max_new_tokens=400,
eos_token_id=terminators,
do_sample=True,
)
result = result[0][input_ids.shape[-1]:]
output = tokenizer.decode(result, skip_special_tokens=True)
print("\n生成結果\n",output)
# 終了時間を記録
end_time = time.time()
# 実行時間を計算
execution_time = end_time - start_time
print(f"Execution time: {execution_time:.2f} seconds")
del input_ids
torch.cuda.empty_cache()
if __name__ == "__main__":
prompt = generate_prompt()
generate_text(prompt)
init.py
# model_setup.py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
# グローバル変数を最初に宣言
model = None
tokenizer = None
model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
#tokenizerとmodelを作成、またはロード
def make_or_load_model_and_tokenizer():
global model, tokenizer
#作成する場合(一回目)
if model is None or tokenizer is None:
print("Loading model and tokenizer")
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=False,
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
quantization_config=bnb_config,
#FlashAttention-2 can only be used when the model’s dtype is fp16 or bf16.
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2"
)
print("Model and tokenizer loaded")
else:
#作成していた場合(1<N回目)
print("Model and tokenizer already loaded")
return model, tokenizer
# モデルとトークナイザーを初期化
model, tokenizer = make_or_load_model_and_tokenizer()