ベクトルデータベース「Weaviate」を試す 5:Generative Searchを使ったRAG

!pip install -U weaviate-client llama-index llama-index-readers-file

from pathlib import Path
import requests
import re

def replace_heading(match):
    level = len(match.group(1))
    return '#' * level + ' ' + match.group(2).strip()

# Wikipediaからのデータ読み込み
wiki_titles = ["オグリキャップ"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = f"# {title}\n\n## 概要\n\n"
    wiki_text += page["extract"]

    wiki_text = re.sub(r"(=+)([^=]+)\1", replace_heading, wiki_text)
    wiki_text = re.sub(r"\t+", "", wiki_text)
    wiki_text = re.sub(r"\n{3,}", "\n\n", wiki_text)
    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.md", "w") as fp:
        fp.write(wiki_text)

from pathlib import Path
import glob
import os

from llama_index.core.node_parser import MarkdownNodeParser
from llama_index.readers.file import FlatReader
from llama_index.core.schema import MetadataMode

files = glob.glob('data/*.md')

docs = []
for f in files:
    doc = FlatReader().load_data(Path(f))
    docs.extend(doc)

parser = MarkdownNodeParser()
nodes = parser.get_nodes_from_documents(docs)

nodes_for_delete = []
sections_for_delete = ["競走成績", "外部リンク", "参考文献", "関連作品"]

for idx, n in enumerate(nodes):
    # メタデータからセクション情報を取り出す。
    metadatas = []
    header_keys = []
    for m in n.metadata:
        if m.startswith("Header"):
            metadatas.append(n.metadata[m])
            header_keys.append(m)
    if len(metadatas) > 0:
        # セクション情報を新たなメタデータに設定
        n.metadata["section"] = metadata_str = " > ".join(metadatas)
        # 古いセクション情報を削除
        for k in header_keys:
            if k.startswith("Header"):
               del n.metadata[k]

    # コンテンツ整形
    contents = n.get_content().split("\n")
    if len(contents) == 1:
        # コンテンツが1つだけ≒セクションタイトルのみの場合は削除対象
        nodes_for_delete.append(idx)
    elif contents[0] in sections_for_delete:
        # 任意のセクションを削除対象
        nodes_for_delete.append(idx)
    else:
        # コンテンツの冒頭にあるセクションタイトル部分、及びそれに続く改行を削除
        content_for_delete = []
        for c_idx, c in enumerate(contents):
            if c in (metadatas):
                content_for_delete.append(c_idx)
            elif c in ["", "\n", None]:
                content_for_delete.append(c_idx)
            else:
                break
        contents = [item for i, item in enumerate(contents) if i not in content_for_delete]

    # 整形したコンテンツでノードを書き換え
    n.set_content("\n".join(contents))

base_nodes = [item for i, item in enumerate(nodes) if i not in nodes_for_delete]

import re
import weaviate.classes as wvc

def text_split(text, max_length=400):
    chunks = re.split(r'([。！？])', text)
    temp_chunk = ""
    final_chunks = []

    for chunk in chunks:
        if len(temp_chunk + chunk) <= max_length:
            temp_chunk += chunk
        else:
            final_chunks.append(temp_chunk)
            temp_chunk = chunk

    if temp_chunk:
        final_chunks.append(temp_chunk)

    return final_chunks

wvc_objs = []
for n in base_nodes:
    content = n.get_content().replace("\n", " ")
    chunks = text_split(content, 400)
    if len(chunks) == 1:
        idx = len(wvc_objs) + 1
        wvc_objs.append(wvc.data.DataObject(
            properties={
                "chunk_id": id, 
                "chapter_title": n.metadata["section"],
                "chunk": chunks[0],
            }
        ))
    else:
        for chunk_idx, chunk in enumerate(chunks, start=1):
            id = len(wvc_objs) + 1
            wvc_objs.append(wvc.data.DataObject(
                properties={
                    "chunk_id": id, 
                    "chapter_title": n.metadata["section"] + f"（{chunk_idx}）",
                    "chunk": chunk,
                }
            ))

print(len(wvc_objs))
print(wvc_objs[0])

107
DataObject(
    properties={
        'chunk_id': 1,
        'chapter_title': 'オグリキャップ > 概要（1）',
        'chunk': 'オグリキャップ（欧字名:Oguri Cap、1985年3月27日 - 2010年7月3日）は、日本の競走馬、種牡馬。 1987年5月に岐阜県の地方競馬・笠松競馬場でデビュー。8連勝、重賞5勝を含む12戦10勝を記録した後、1988年1月に中央競馬へ移籍し、重賞12勝（うちGI4勝）を記録した。1988年度のJRA賞最優秀4歳牡馬、1989年度のJRA賞特別賞、1990年度のJRA賞最優秀5歳以上牡馬および年度代表馬。1991年、JRA顕彰馬に選出。愛称は「オグリ」「芦毛の怪物」など多数。 中央競馬時代はスーパークリーク、イナリワンの二頭とともに「平成三強」と総称され、自身と騎手である武豊の活躍を中心として起こった第二次競馬ブーム期において、第一次競馬ブームの立役者とされるハイセイコーに比肩するとも評される高い人気を得た。'
    },
    uuid=None,
    vector=None,
    references=None
)

import weaviate
import os
import requests
import json
from google.colab import userdata

client = weaviate.connect_to_wcs(
    cluster_url=userdata.get('WEAVIATE_CLUSTER_URL'),
    auth_credentials=weaviate.auth.AuthApiKey(userdata.get('WEAVIATE_API_KEY')),
    headers={
        "X-OpenAI-Api-Key": userdata.get('OPENAI_API_KEY')
    }
)

client.collections.delete(name="OguriCap")

oguricap = client.collections.create(
    name="OguriCap",
    vectorizer_config=wvc.config.Configure.Vectorizer.text2vec_openai(
        model="text-embedding-3-small",
        vectorize_collection_name=False
    ),
    generative_config=wvc.config.Configure.Generative.openai(
        model="gpt-3.5-turbo"
    ),
    properties=[
        wvc.config.Property(
            name="chunk_id",
            data_type=wvc.config.DataType.INT,
        ),
        wvc.config.Property(
            name="chapter_title",
            data_type=wvc.config.DataType.TEXT,
        ),
        wvc.config.Property(
            name="chunk",
            data_type=wvc.config.DataType.TEXT,
        ),
    ]
)

oguricap.data.insert_many(wvc_objs)

response = oguricap.aggregate.over_all(total_count=True)
print(response.total_count)

107

response = oguricap.generate.fetch_objects(
    limit=2,
    single_prompt="次の文章を下に、俳句を1句読み上げて下さい: ===== {chunk}"
)

for o in response.objects:
    print(f"\n===== Object index: [{o.properties['chunk_id']}] =====")
    print(o.generated)

===== Object index: [78] =====
春の風や　オグリキャップの脚　軽やかに

===== Object index: [85] =====
食べる馬　オグリキャップ　食欲旺盛

response = oguricap.generate.fetch_objects(
    limit=2,
    single_prompt="要約して下さい。"
)

for o in response.objects:
    print(f"\n===== Object index: [{o.properties['chunk_id']}] =====")
    print(o.generated)

WeaviateQueryError: Query call with protocol GRPC search failed with message explorer: list class: extend: extend generate: Prompt does not contain any properties. Use {PROPERTY_NAME} in the prompt to instuct Weaviate which data to use.

response = oguricap.generate.fetch_objects(
    limit=5,
    grouped_task="箇条書きで要約して下さい。"
)

for o in response.objects:
    print(f"===== Object index: [{o.properties['chunk_id']}] =====")
    print(f"chapter_title: [{o.properties['chapter_title']}]")
    print(f"chunk: [{o.properties['chunk']}]")
    print()

print("回答: ", response.generated)

===== Object index: [78] =====
chapter_title: [オグリキャップ > 特徴・評価 > 総合的な評価（1）]
chunk: [鷲見昌勇はオグリキャップが3歳の時点で「五十年に一頭」「もうあんなにすごい馬は笠松からは出ないかもしれない」と述べている。安藤勝己は初めて調教のためにオグリキャップに騎乗したとき、厩務員の川瀬に「どえらい馬だね。来年は間違いなく東海ダービーを取れる」と言った。安藤のオグリキャップに対する評価は高く、3歳の時点で既に「オグリキャップを負かすとすればフェートノーザンかワカオライデンのどちらか」と考えていた。 河内洋は初めて騎乗したペガサスステークスについて「とても長くいい脚を使えたし、これはかなり走りそうだと思ったよ。距離の融通も利きそうだったしね」と回顧し、ニュージーランドトロフィー4歳ステークスのレース後に古馬との比較について問われた際、キャリアの違いはあったものの「この馬も相当の器だよ」とコメントした。]

===== Object index: [85] =====
chapter_title: [オグリキャップ > 特徴・評価 > 競走馬名および愛称・呼称]
chunk: [競走馬名「オグリキャップ」の由来は、馬主の小栗が使用していた冠名「オグリ」に父ダンシングキャップの馬名の一部「キャップ」を加えたものである。 同馬の愛称としては「オグリ」が一般的だが、女性ファンの中には「オグリちゃん」、「オグリン」と呼ぶファンも存在し、その他「怪物」「新怪物」「白い怪物」「芦毛の怪物」と呼ばれた。またオグリキャップは前述のように生来食欲が旺盛で、「食べる競走馬」とも呼ばれた。]

===== Object index: [17] =====
chapter_title: [オグリキャップ > 競走馬時代 > 中央競馬時代 > 4歳（1988年） > 競走内容（4）]
chunk: [ 続く高松宮杯では、中央競馬移籍後初の古馬との対戦、特に重賞優勝馬でありこの年の宝塚記念で4着となったランドヒリュウとの対戦にファンの注目が集まった。レースではランドヒリュウが先頭に立って逃げたのに対してオグリキャップは序盤は4番手に位置して第3コーナーから前方への進出を開始する。第4コーナーで2番手に立つと直線でランドヒリュウをかわし、中京競馬場芝2000mのコースレコードを記録して優勝した。この勝利により、地方競馬からの移籍馬による重賞連勝記録である5連勝を達成した。 高松宮杯のレース後、陣営は秋シーズンのオグリキャップのローテーションを検討し、毎日王冠を経て天皇賞（秋）でGIに初出走することを決定した。毎日王冠までは避暑 を行わず、栗東トレーニングセンターで調整を行い、8月下旬から本格的な調教を開始。9月末に東京競馬場に移送された。]

===== Object index: [18] =====
chapter_title: [オグリキャップ > 競走馬時代 > 中央競馬時代 > 4歳（1988年） > 競走内容（5）]
chunk: [ 毎日王冠では終始後方からレースを進め、第3コーナーからまくりをかけて優勝した。この勝利により、当時のJRA重賞連勝記録である6連勝を達成した（メジロラモーヌと並ぶタイ記録）。当時競馬評論家として活動していた大橋巨泉は、オグリキャップのレース内容について「毎日王冠で古馬の一線級を相手に、スローペースを後方から大外廻って、一気に差し切るなどという芸当は、今まで見たことがない」「どうやらオグリキャップは本当のホンモノの怪物らしい」と評した。毎日王冠の後、オグリキャップはそのまま東京競馬場に留まって調整を続けた（レースに関する詳細については第39回毎日王冠を参照）。 続く天皇賞（秋）では、前年秋から7連勝中であった古馬のタマモクロスを凌いで1番人気に支持された。]

===== Object index: [97] =====
chapter_title: [オグリキャップ > 人気 > 第35回有馬記念優勝後の人気（2）]
chunk: [武豊は1998年に受けたインタビューにおいて、同レースが「奇蹟」などと言われていることについて、「こんな言い方は失礼かもしれないけど、オグリよりも、あの時（翌1991年の有馬記念）のダイユウサクの脚のほうが『奇蹟』でしょう」と述べている。 第35回有馬記念はNHKとフジテレビが生放送し、ビデオリサーチの発表によると視聴率はそれぞれ11.7%と9.6%だった。前年の有馬記念はNHKが16.2%、フジテレビが10.4%を記録していたが、2局合わせた番組占拠率は50.3%を記録した。]

回答:  1. 鷲見昌勇と安藤勝己はオグリキャップの才能を高く評価しており、河内洋もその器の大きさを認めている。
2. オグリキャップの名前の由来は「オグリ」＋「キャップ」で、愛称は「オグリ」だが他にも様々な呼び方がある。
3. オグリキャップは移籍後の高松宮杯でコースレコードを樹立し、5連勝を達成。その後も毎日王冠で6連勝を果たす。
4. 有馬記念優勝後のオグリキャップは人気が高まり、武豊はその偉業を称賛している。

query = "オグリキャップに騎乗した騎手を全員リストアップして下さい"

response = oguricap.generate.near_text(
    limit=5,
    grouped_task=f"事前知識を使わずに、与えられたコンテキストだけを使って、以下に質問に回答して下さい: {query}",
    query=query
)

for o in response.objects:
    print(f"===== Object index: [{o.properties['chunk_id']}] =====")
    print(f"chapter_title: [{o.properties['chapter_title']}]")
    print(f"chunk: [{o.properties['chunk']}]")
    print()

print("回答: ", response.generated)

===== Object index: [98] =====
chapter_title: [オグリキャップ > 騎手（1）]
chunk: [オグリキャップの騎手は何度も交替した。以下、オグリキャップに騎乗した主な騎手と騎乗した経緯について記述する。  安藤勝己 デビュー当初のオグリキャップには高橋一成（鷲見厩舎の所属騎手）と青木達彦が騎乗していたが、6戦目のレースではいずれも騎乗することができなかったため、当時笠松競馬場のリーディングジョッキーであった安藤が騎乗し、以降は安藤がオグリキャップの主戦騎手を務めた。安藤はオグリキャップに騎乗した中でもっとも思い出に残るレースとして、楽に勝たせようと思い早めに先頭に立った結果オグリキャップが気を抜いてマーチトウショウとの競り合いになり、目一杯に追う羽目になった7戦目のジュニアクラウンを挙げている。1990年のジャパンカップで川崎競馬場所属の騎手河津裕昭がイブンベイに騎乗すると聞き、佐橋に対しオグリキャップへの騎乗を申し入れたが実現しなかった。]

===== Object index: [100] =====
chapter_title: [オグリキャップ > 騎手（3）]
chunk: [ 岡部幸雄 1988年の有馬記念で佐橋の意向から騎乗依頼を受け「一回だけ」という条件付きで依頼を引き受けた。後に佐橋からオグリキャップを購入した近藤俊典が騎乗依頼を出したが、了解を得ることはできなかった。その一因は、クリーンなイメージを大切にする岡部が佐橋の脱税問題が取りざたされたオグリキャップへの騎乗を嫌ったことにあるとされている。調教師の高松邦男は、オグリキャップに騎乗した騎手の中で岡部がもっとも馬にフィットした乗り方をしたと評した。岡部は有馬記念前に騎乗依頼を受けた際、「西（栗東）の馬はよくわからないから」と一度は婉曲に断っていたが、オグリキャップが引退して一年余り経ったときに、「僕はオグリキャップには乗ってみたかったんですよ。ひと口に強い馬といってもいろいろ癖がありますから。どんな馬か知りたかったですしね」とオグリキャップへの騎乗に興味を持っていたことを明かしている。]

===== Object index: [105] =====
chapter_title: [オグリキャップ > 騎手（8）]
chunk: [ 増沢末夫 1990年の天皇賞（秋）およびジャパンカップに騎乗した。当時の馬主である近藤俊典は、増沢が非常に可愛がられていた馬主・近藤ハツ子の甥であり、増沢が若手騎手の頃から面識があった。前述のように、安藤勝己、青木達彦、渡瀬夏彦は増沢とオグリキャップとの相性は良くなかったという見解を示している。一方、瀬戸口は「増沢騎手には本当に気の毒な思いをさせました。済まなかったと思います。あれほどの騎手に、オグリがいちばん体調がよくない時に乗ってもらったんですから」と、増沢を庇う言葉を残している。]

===== Object index: [99] =====
chapter_title: [オグリキャップ > 騎手（2）]
chunk: [安藤は後にJRAへ移籍したが、「オグリキャップがいたから中央競馬が地方馬にGI開放とかそういう流れになっていったんじゃないかと思いますし、（自身の中央移籍への道を作ってくれたのも）オグリキャップのお陰だと思っています」と述べ、自身にとってのオグリキャップの存在についても「自分の未来を切り開いてくれた馬だと思います。とても感謝しています」と述べている。 河内洋 オグリキャップが中央競馬へ移籍した当初の主戦騎手。当時の馬主の佐橋の強い意向によりペガサスステークスに騎乗し、同年のジャパンカップまで主戦騎手を務めた。同年の有馬記念では、佐橋の意向により岡部幸雄に騎乗依頼が出され、それ以降、オグリキャップに騎乗することはなかった。河内は自身が騎乗したレースで最も印象に残っているレースとしてニュージーランドT4歳Sを挙げている。]

===== Object index: [101] =====
chapter_title: [オグリキャップ > 騎手（4）]
chunk: [ 南井克巳 1988年の京都4歳特別に、河内洋の代役として騎乗した。翌1989年には前述のように岡部が近藤からの騎乗依頼を断った後で瀬戸口から騎乗依頼を受け、主戦騎手を務めた。1990年はバンブービギンに騎乗することを決断し、自ら降板を申し出た。1989年の第34回有馬記念における騎乗について野平祐二と岡部は南井の騎乗ミスを指摘した。南井は有馬記念の騎乗について、オグリキャップの調子が悪くいつもの末脚を発揮することが難しいため、好位置の楽な競馬で気力を取り戻すことを期待したと説明している。 武豊 1990年にアメリカ遠征が決定した際、武豊が鞍上を務めることが決まったことで陣営は第40回安田記念にも騎乗を依頼し、同レースで騎乗することとなった。]

回答:  安藤勝己、岡部幸雄、増沢末夫、河内洋、南井克巳、武豊

response = oguricap.generate.fetch_objects()
response.objects[0].properties