AGIと未来のAI
Rerankers と Two-Stage Retrieval の力 – Retrieval Augmented Generation におけるその重要性
このコードスニペットは、jina-colbert-v1-enモデルを使用してドキュメントコレクションをインデックス化し、その長いコンテキストを効率的に処理する能力を示しています。
Two-Stage Retrieval with Rerankersの実装
Two-Stage RetrievalとRerankersの原理を理解したので、RAGシステム内での実践的な実装について探りましょう。人気のライブラリとフレームワークを使用して、これらの技術の統合を示します。
環境の設定
コードに取り組む前に、開発環境を設定しましょう。Pythonと人気のNLPライブラリ、Hugging Face Transformers、Sentence Transformers、LanceDBを使用します。
# 必要なライブラリをインストール !pip install datasets huggingface_hub sentence_transformers lancedb
データの準備
デモンストレーションのために、Hugging Face Datasetsから「ai-arxiv-chunked」データセットを使用します。このデータセットには、機械学習、自然言語処理、そして大規模言語モデルに関する400以上のarXiv論文が含まれています。
from datasets import load_dataset <p>dataset = load_dataset("jamescalam/ai-arxiv-chunked", split="train")</p> &lt;pre&gt;
次に、データを前処理し、小さなチャンクに分割して効率的な検索と処理を促進します。
from transformers import AutoTokenizer <p>tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")</p> <p>def chunk_text(text, chunk_size=512, overlap=64): tokens = tokenizer.encode(text, return_tensors="pt", truncation=True) chunks = tokens.split(chunk_size - overlap) texts = [tokenizer.decode(chunk) for chunk in chunks] return texts</p> <p>chunked_data = [] for doc in dataset: text = doc["chunk"] chunked_texts = chunk_text(text) chunked_data.extend(chunked_texts)</p>
初期検索ステージでは、Sentence Transformerモデルを使用してドキュメントとクエリを密なベクトル表現にエンコードし、次にLanceDBなどのベクトルデータベースを使用して近似最近傍探索を実行します。
from sentence_transformers import SentenceTransformer from lancedb import lancedb <p># Sentence Transformerモデルをロード model = SentenceTransformer('all-MiniLM-L6-v2')</p> <p># LanceDBベクトルストアを作成 db = lancedb.lancedb('/path/to/store') db.create_collection('docs', vector_dimension=model.get_sentence_embedding_dimension())</p> <p># ドキュメントをインデックス化 for text in chunked_data: vector = model.encode(text).tolist() db.insert_document('docs', vector, text)</p> <p>from sentence_transformers import SentenceTransformer from lancedb import lancedb</p> <p># Sentence Transformerモデルをロード model = SentenceTransformer('all-MiniLM-L6-v2')</p> <p># LanceDBベクトルストアを作成 db = lancedb.lancedb('/path/to/store') db.create_collection('docs', vector_dimension=model.get_sentence_embedding_dimension())</p> <p># ドキュメントをインデックス化 for text in chunked_data: vector = model.encode(text).tolist() db.insert_document('docs', vector, text)
ドキュメントがインデックス化されたら、クエリベクトルに最も近いものを検索することで初期検索を実行できます。
from transformers import AutoTokenizer <p>tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")</p> <p>def chunk_text(text, chunk_size=512, overlap=64): tokens = tokenizer.encode(text, return_tensors="pt", truncation=True) chunks = tokens.split(chunk_size - overlap) texts = [tokenizer.decode(chunk) for chunk in chunks] return texts</p> <p>chunked_data = [] for doc in dataset: text = doc["chunk"] chunked_texts = chunk_text(text) chunked_data.extend(chunked_texts)
再ランキング
初期検索の後、ColBERT再ランカーを使用して、検索されたドキュメントをクエリへの関連性に基づいて再ランキングします。
from lancedb.rerankers import ColbertReranker reranker = ColbertReranker() <p># 初期ドキュメントを再ランキング reranked_docs = reranker.rerank(query, initial_docs)</p>
再ランキングされたドキュメントのリストreranked_docsには、クエリへの関連性に基づいて再ランキングされたドキュメントが含まれます。
増強と生成
関連性の高いドキュメントを再ランキングした後、RAGパイプラインの増強と生成ステージに進みます。最終的な応答を生成するために、Hugging Face Transformersライブラリから言語モデルを使用します。
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM <p>tokenizer = AutoTokenizer.from_pretrained("t5-base") model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")</p> <p># クエリを再ランキングされたドキュメントと増強 augmented_query = query + " " + " ".join(reranked_docs[:3])</p> <p># 言語モデルから応答を生成 input_ids = tokenizer.encode(augmented_query, return_tensors="pt") output_ids = model.generate(input_ids, max_length=500) response = tokenizer.decode(output_ids[0], skip_special_tokens=True)</p> print(response)














