超軽量なColBERTモデル「LFM2-ColBERT-350M」を試す

pip install pylate

Successfully installed fast-plaid-1.2.4.280 fastkmeans-0.5.0 maturin-1.9.6 ninja-1.11.1.4 pylate-1.3.4 sentence-transformers-5.1.1 setuptools-80.9.0 sqlitedict-2.1.0 transformers-4.56.2 ujson-5.10.0 voyager-2.1.0

from pylate import indexes, models, retrieve

# Step1: ColBERTモデルをロード
model = models.ColBERT(
    model_name_or_path="LiquidAI/LFM2-ColBERT-350M",
)
model.tokenizer.pad_token = model.tokenizer.eos_token

# Step2: PLAIDインデックスを初期化
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # 既存のインデックスがあれば上書き
)

# Step3: ドキュメントをインデックスに
documents_ids = ["1", "2", "3", "4", "5", "6", "7", "8"]
documents = [
    "パリはフランスの首都であり、同国最大の都市である。フランス北部を流れるセーヌ川流域に位置し、政治・経済・文化の中心地として機能している。",
    "東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。",
    "キャンベラはオーストラリアの首都である。シドニーやメルボルンとは異なり、首都として特別に計画・建設された都市で、オーストラリア南東部のオーストラリア首都特別地域に位置している。",
    "ベルリンはドイツの首都であり、同国最大の都市である。1990年のベルリンの壁崩壊後に再統一され、現在ではヨーロッパにおける主要な文化・政治の中心地となっている。",
    "オタワはカナダの首都である。オンタリオ州東部、オタワ川の南岸に位置し、1857年にヴィクトリア女王によって首都として選定された。",
    "ブラジリアはブラジルの首都である。1960年に特別に計画・建設された都市で、都市計画家ルシオ・コスタと建築家オスカー・ニーマイヤーによって設計され、その特徴的なモダニズム建築で知られている。",
    "カイロはエジプトの首都であり、アフリカ大陸でも有数の大都市である。ナイル川デルタ地帯に位置し、数千年にわたって文明の中心地としての役割を果たしてきた。",
    "ウェリントンはニュージーランドの首都である。北島南西端に位置し、世界で最も風の強い都市の一つとして知られている。",
]

# ドキュメントをエンコード
documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,  # ドキュメントの場合は False
    show_progress_bar=True,
)

# ドキュメントのIDと埋め込みをPLAIDインデックスに追加
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

# Step4: 検索
retriever = retrieve.ColBERT(index=index)

queries = [
    "日本の首都は？",
    "フランスの首都は",
    "Which city is Japan's capital?",
    "What is the capital of France?"
]

queries_embeddings = model.encode(
    queries,
    batch_size=32,
    is_query=True, # クエリの場合は False
    show_progress_bar=True,
)

scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=3,
)

print(scores)

[
    [
        {'id': '2', 'score': 30.41943359375},
        {'id': '4', 'score': 29.88232421875},
        {'id': '5', 'score': 29.86083984375}
    ],
    [
        {'id': '1', 'score': 30.34130859375},
        {'id': '2', 'score': 30.01416015625},
        {'id': '5', 'score': 29.93505859375}
    ],
    [
        {'id': '2', 'score': 30.22705078125},
        {'id': '6', 'score': 30.10791015625},
        {'id': '1', 'score': 30.05712890625}
    ],
    [
        {'id': '1', 'score': 30.462890625},
        {'id': '6', 'score': 30.36328125},
        {'id': '4', 'score': 30.30322265625}
    ]
]

for i, query_scores in enumerate(scores):
    print(f"クエリ: {queries[i]}")
    for score_info in query_scores:
        doc_id = score_info['id']
        score = score_info['score']
        doc_index = documents_ids.index(doc_id)
        document_text = documents[doc_index]
        print(f"  ID: {doc_id}, Score: {score:.2f}, Document: {document_text}")
    print("-" * 20)

クエリ: 日本の首都は？
  ID: 2, Score: 30.42, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 4, Score: 29.88, Document: ベルリンはドイツの首都であり、同国最大の都市である。1990年のベルリンの壁崩壊後に再統一され、現在ではヨーロッパにおける主要な文化・政治の中心地となっている。
  ID: 5, Score: 29.86, Document: オタワはカナダの首都である。オンタリオ州東部、オタワ川の南岸に位置し、1857年にヴィクトリア女王によって首都として選定された。
--------------------
クエリ: フランスの首都は
  ID: 1, Score: 30.34, Document: パリはフランスの首都であり、同国最大の都市である。フランス北部を流れるセーヌ川流域に位置し、政治・経済・文化の中心地として機能している。
  ID: 2, Score: 30.01, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 5, Score: 29.94, Document: オタワはカナダの首都である。オンタリオ州東部、オタワ川の南岸に位置し、1857年にヴィクトリア女王によって首都として選定された。
--------------------
クエリ: Which city is Japan's capital?
  ID: 2, Score: 30.23, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 6, Score: 30.11, Document: ブラジリアはブラジルの首都である。1960年に特別に計画・建設された都市で、都市計画家ルシオ・コスタと建築家オスカー・ニーマイヤーによって設計され、その特徴的なモダニズム建築で知られている。
  ID: 1, Score: 30.06, Document: パリはフランスの首都であり、同国最大の都市である。フランス北部を流れるセーヌ川流域に位置し、政治・経済・文化の中心地として機能している。
--------------------
クエリ: What is the capital of France?
  ID: 1, Score: 30.46, Document: パリはフランスの首都であり、同国最大の都市である。フランス北部を流れるセーヌ川流域に位置し、政治・経済・文化の中心地として機能している。
  ID: 6, Score: 30.36, Document: ブラジリアはブラジルの首都である。1960年に特別に計画・建設された都市で、都市計画家ルシオ・コスタと建築家オスカー・ニーマイヤーによって設計され、その特徴的なモダニズム建築で知られている。
  ID: 4, Score: 30.30, Document: ベルリンはドイツの首都であり、同国最大の都市である。1990年のベルリンの壁崩壊後に再統一され、現在ではヨーロッパにおける主要な文化・政治の中心地となっている。
--------------------

!sudo apt update && sudo apt install -y tree
!tree pylate-index

pylate-index
└── index
    ├── documents_ids_to_plaid_ids.sqlite
    ├── fast_plaid_index
    │   ├── 0.codes.npy
    │   ├── 0.metadata.json
    │   ├── 0.residuals.npy
    │   ├── avg_residual.npy
    │   ├── bucket_cutoffs.npy
    │   ├── bucket_weights.npy
    │   ├── centroids.npy
    │   ├── doclens.0.json
    │   ├── ivf_lengths.npy
    │   ├── ivf.npy
    │   ├── metadata.json
    │   └── plan.json
    └── plaid_ids_to_documents_ids.sqlite

2 directories, 14 files

index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
)

from pylate import indexes, models, retrieve

# ColBERTモデルをロード
model = models.ColBERT(
    model_name_or_path="LiquidAI/LFM2-ColBERT-350M",
)
model.tokenizer.pad_token = model.tokenizer.eos_token

# 既存のPLAIDインデックスをロード
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
)

# 表示用にドキュメントを定義（インデックスには既に登録済み）
documents_ids = ["1", "2", "3", "4", "5", "6", "7", "8"]
documents = [
    "パリはフランスの首都であり、同国最大の都市である。フランス北部を流れるセーヌ川流域に位置し、政治・経済・文化の中心地として機能している。",
    "東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。",
    "キャンベラはオーストラリアの首都である。シドニーやメルボルンとは異なり、首都として特別に計画・建設された都市で、オーストラリア南東部のオーストラリア首都特別地域に位置している。",
    "ベルリンはドイツの首都であり、同国最大の都市である。1990年のベルリンの壁崩壊後に再統一され、現在ではヨーロッパにおける主要な文化・政治の中心地となっている。",
    "オタワはカナダの首都である。オンタリオ州東部、オタワ川の南岸に位置し、1857年にヴィクトリア女王によって首都として選定された。",
    "ブラジリアはブラジルの首都である。1960年に特別に計画・建設された都市で、都市計画家ルシオ・コスタと建築家オスカー・ニーマイヤーによって設計され、その特徴的なモダニズム建築で知られている。",
    "カイロはエジプトの首都であり、アフリカ大陸でも有数の大都市である。ナイル川デルタ地帯に位置し、数千年にわたって文明の中心地としての役割を果たしてきた。",
    "ウェリントンはニュージーランドの首都である。北島南西端に位置し、世界で最も風の強い都市の一つとして知られている。",
]

# 検索
retriever = retrieve.ColBERT(index=index)

queries = ["日本の首都は？"]

queries_embeddings = model.encode(
    queries,
    batch_size=32,
    is_query=True,
    show_progress_bar=True,
)

scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=3,
)

for i, query_scores in enumerate(scores):
    print(f"クエリ: {queries[i]}")
    for score_info in query_scores:
        doc_id = score_info['id']
        score = score_info['score']
        doc_index = documents_ids.index(doc_id)
        document_text = documents[doc_index]
        print(f"  ID: {doc_id}, Score: {score:.2f}, Document: {document_text}")
    print("-" * 20)

クエリ: 日本の首都は？
  ID: 2, Score: 30.42, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 4, Score: 29.88, Document: ベルリンはドイツの首都であり、同国最大の都市である。1990年のベルリンの壁崩壊後に再統一され、現在ではヨーロッパにおける主要な文化・政治の中心地となっている。
  ID: 5, Score: 29.86, Document: オタワはカナダの首都である。オンタリオ州東部、オタワ川の南岸に位置し、1857年にヴィクトリア女王によって首都として選定された。
--------------------

from pylate import indexes, models, retrieve

model = models.ColBERT(
    model_name_or_path="LiquidAI/LFM2-ColBERT-350M",
)
model.tokenizer.pad_token = model.tokenizer.eos_token

index = indexes.PLAID(
    index_folder="pylate-crosslingual-index",
    index_name="index",
    override=True,
)

documents_ids = ["1", "2", "3", "4", "5", "6", "7", "8"]
documents = [
    # 1. パリ: フランス語
    "Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.",
    # 2. 東京: 日本語
    "東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。",
    # 3. キャンベラ: 英語
    "Canberra is the capital city of Australia. Unlike Sydney or Melbourne, Canberra was purpose-built as the capital and is located in the Australian Capital Territory in southeastern Australia.",
    # 4. ベルリン: ドイツ語
    "Berlin ist die Hauptstadt und größte Stadt Deutschlands. Nach der Wiedervereinigung im Jahr 1990 nach dem Fall der Berliner Mauer ist sie heute ein wichtiges kulturelles und politisches Zentrum Europas.",
    # 5. カイロ: アラビア語
    "القاهرة هي عاصمة مصر وإحدى أكبر المدن في أفريقيا. تقع بالقرب من دلتا النيل، وكانت مركزًا للحضارة لآلاف السنين.",
    # 6. マドリード: スペイン語
    "Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.",
    # 7. 北京: 中国語（簡体字）
    "北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。",
    # 8. ソウル: 韓国語
    "서울은 대한민국의 수도이자 최대 도시로, 정치·경제·문화의 중심지 역할을 하고 있습니다.",
]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,
    show_progress_bar=True,
)

index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

retriever = retrieve.ColBERT(index=index)

queries = [
    "Quelle est la capitale du Japon ?",     # フランス語 / 「日本の首都はどこですか？」
    "オーストラリアの首都はどこですか？",         # 日本語 / 「オーストラリアの首都はどこですか？」
    "What is the capital of Germany?",       # 英語 / 「ドイツの首都はどこですか？」
    "Was ist die Hauptstadt von Ägypten?",   # ドイツ語 / 「エジプトの首都はどこですか？」
    "ما عاصمة إسبانيا؟",                      # アラビア語 / 「スペインの首都はどこですか？」
    "¿Cuál es la capital de China?",         # スペイン語 / 「中国の首都はどこですか？」
    "韩国的首都是哪座城市？",                   # 中国語 / 「韓国の首都はどの都市ですか？」
    "프랑스의 수도는 어디예요?",                # 韓国語 / 「フランスの首都はどこですか？」
]

queries_embeddings = model.encode(
    queries,
    batch_size=32,
    is_query=True,
    show_progress_bar=True,
)

scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=3,
)

for i, query_scores in enumerate(scores):
    print(f"クエリ: {queries[i]}")
    for score_info in query_scores:
        doc_id = score_info["id"]
        score = score_info["score"]
        doc_index = documents_ids.index(doc_id)
        document_text = documents[doc_index]
        print(f"  ID: {doc_id}, Score: {score:.2f}, Document: {document_text}")
    print("-" * 20)

クエリ: Quelle est la capitale du Japon ?
  ID: 1, Score: 29.85, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 2, Score: 29.84, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 6, Score: 29.74, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
--------------------
クエリ: オーストラリアの首都はどこですか？
  ID: 2, Score: 29.56, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 7, Score: 29.40, Document: 北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。
  ID: 3, Score: 29.29, Document: Canberra is the capital city of Australia. Unlike Sydney or Melbourne, Canberra was purpose-built as the capital and is located in the Australian Capital Territory in southeastern Australia.
--------------------
クエリ: What is the capital of Germany?
  ID: 4, Score: 30.55, Document: Berlin ist die Hauptstadt und größte Stadt Deutschlands. Nach der Wiedervereinigung im Jahr 1990 nach dem Fall der Berliner Mauer ist sie heute ein wichtiges kulturelles und politisches Zentrum Europas.
  ID: 6, Score: 30.41, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
  ID: 1, Score: 30.40, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
--------------------
クエリ: Was ist die Hauptstadt von Ägypten?
  ID: 4, Score: 30.09, Document: Berlin ist die Hauptstadt und größte Stadt Deutschlands. Nach der Wiedervereinigung im Jahr 1990 nach dem Fall der Berliner Mauer ist sie heute ein wichtiges kulturelles und politisches Zentrum Europas.
  ID: 1, Score: 29.93, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 6, Score: 29.91, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
--------------------
クエリ: ما عاصمة إسبانيا؟
  ID: 6, Score: 30.13, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
  ID: 1, Score: 29.91, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 5, Score: 29.91, Document: القاهرة هي عاصمة مصر وإحدى أكبر المدن في أفريقيا. تقع بالقرب من دلتا النيل، وكانت مركزًا للحضارة لآلاف السنين.
--------------------
クエリ: ¿Cuál es la capital de China?
  ID: 7, Score: 30.04, Document: 北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。
  ID: 1, Score: 30.03, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 6, Score: 30.03, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
--------------------
クエリ: 韩国的首都是哪座城市？
  ID: 7, Score: 29.37, Document: 北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。
  ID: 2, Score: 29.21, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 8, Score: 29.14, Document: 서울은 대한민국의 수도이자 최대 도시로, 정치·경제·문화의 중심지 역할을 하고 있습니다.
--------------------
クエリ: 프랑스의 수도는 어디예요?
  ID: 1, Score: 29.29, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 8, Score: 29.25, Document: 서울은 대한민국의 수도이자 최대 도시로, 정치·경제·문화의 중심지 역할을 하고 있습니다.
  ID: 2, Score: 29.20, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
--------------------

from pylate import rank

queries = [
    "query A",
    "query B",
]

documents = [
    ["document A", "document B"],
    ["document 1", "document C", "document B"],
]

documents_ids = [
    [1, 2],
    [1, 3, 2],
]

queries_embeddings = model.encode(
    queries,
    is_query=True,
)

documents_embeddings = model.encode(
    documents,
    is_query=False,
)

reranked_documents = rank.rerank(
    documents_ids=documents_ids,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings,
)

from pylate import indexes, models, retrieve, rank

model = models.ColBERT(
    model_name_or_path="LiquidAI/LFM2-ColBERT-350M",
)
model.tokenizer.pad_token = model.tokenizer.eos_token

documents_ids = ["1", "2", "3", "4", "5", "6", "7", "8"]
documents = [
    # 1. パリ: フランス語
    "Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.",
    # 2. 東京: 日本語
    "東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。",
    # 3. キャンベラ: 英語
    "Canberra is the capital city of Australia. Unlike Sydney or Melbourne, Canberra was purpose-built as the capital and is located in the Australian Capital Territory in southeastern Australia.",
    # 4. ベルリン: ドイツ語
    "Berlin ist die Hauptstadt und größte Stadt Deutschlands. Nach der Wiedervereinigung im Jahr 1990 nach dem Fall der Berliner Mauer ist sie heute ein wichtiges kulturelles und politisches Zentrum Europas.",
    # 5. カイロ: アラビア語
    "القاهرة هي عاصمة مصر وإحدى أكبر المدن في أفريقيا. تقع بالقرب من دلتا النيل、 وكانت مركزًا للحضارة لآلاف السنين.",
    # 6. マドリード: スペイン語
    "Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.",
    # 7. 北京: 中国語（簡体字）
    "北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。",
    # 8. ソウル: 韓国語
    "서울은 대한민국의 수도이자 최대 도시로, 정치·경제·문화의 중심지 역할을 하고 있습니다.",
]

# ドキュメントの埋め込みを一度だけ生成
documents_embeddings_single = model.encode(
    documents,
    is_query=False,
    show_progress_bar=True,
)

queries = [
    "Quelle est la capitale du Japon ?",     # フランス語 / 「日本の首都はどこですか？」
    "オーストラリアの首都はどこですか？",         # 日本語 / 「オーストラリアの首都はどこですか？」
    "What is the capital of Germany?",       # 英語 / 「ドイツの首都はどこですか？」
    "Was ist die Hauptstadt von Ägypten?",   # ドイツ語 / 「エジプトの首都はどこですか？」
    "ما عاصمة إسبانيا؟",                      # アラビア語 / 「スペインの首都はどこですか？」
    "¿Cuál es la capital de China?",         # スペイン語 / 「中国の首都はどこですか？」
    "韓国的首都是哪座城市？",                   # 中国語 / 「韓国の首都はどの都市ですか？」
    "프랑스의 수도는 어디예요?",                # 韓国語 / 「フランスの首都はどこですか？」
]

queries_embeddings = model.encode(
    queries,
    is_query=True,
    show_progress_bar=True,
)

# 各クエリに対して同じドキュメントIDと埋め込みのリストを参照するようにリストを作成
documents_ids_for_rerank = [documents_ids] * len(queries)
documents_embeddings_for_rerank = [documents_embeddings_single] * len(queries)

# リランキングを実行
reranked_documents = rank.rerank(
    documents_ids=documents_ids_for_rerank,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings_for_rerank,
)

# リランキング結果から上位3件を取得
reranked_documents_top_k = [result[:3] for result in reranked_documents]

# 結果を表示
for i, query_scores in enumerate(reranked_documents_top_k):
    print(f"クエリ: {queries[i]}")
    for score_info in query_scores:
        doc_id = score_info["id"]
        score = score_info["score"]
        # ドキュメントIDから元のdocumentsリストのインデックスを取得
        doc_index = documents_ids.index(doc_id)
        document_text = documents[doc_index]
        print(f"  ID: {doc_id}, Score: {score:.2f}, Document: {document_text}")
    print("-" * 20)

クエリ: Quelle est la capitale du Japon ?
  ID: 1, Score: 29.84, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 2, Score: 29.83, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 6, Score: 29.72, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
--------------------
クエリ: オーストラリアの首都はどこですか？
  ID: 2, Score: 29.56, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 7, Score: 29.39, Document: 北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。
  ID: 3, Score: 29.29, Document: Canberra is the capital city of Australia. Unlike Sydney or Melbourne, Canberra was purpose-built as the capital and is located in the Australian Capital Territory in southeastern Australia.
--------------------
クエリ: What is the capital of Germany?
  ID: 4, Score: 30.54, Document: Berlin ist die Hauptstadt und größte Stadt Deutschlands. Nach der Wiedervereinigung im Jahr 1990 nach dem Fall der Berliner Mauer ist sie heute ein wichtiges kulturelles und politisches Zentrum Europas.
  ID: 6, Score: 30.40, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
  ID: 1, Score: 30.39, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
--------------------
クエリ: Was ist die Hauptstadt von Ägypten?
  ID: 4, Score: 30.08, Document: Berlin ist die Hauptstadt und größte Stadt Deutschlands. Nach der Wiedervereinigung im Jahr 1990 nach dem Fall der Berliner Mauer ist sie heute ein wichtiges kulturelles und politisches Zentrum Europas.
  ID: 1, Score: 29.93, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 6, Score: 29.90, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
--------------------
クエリ: ما عاصمة إسبانيا؟
  ID: 6, Score: 30.12, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
  ID: 5, Score: 29.90, Document: القاهرة هي عاصمة مصر وإحدى أكبر المدن في أفريقيا. تقع بالقرب من دلتا النيل、 وكانت مركزًا للحضارة لآلاف السنين.
  ID: 1, Score: 29.90, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
--------------------
クエリ: ¿Cuál es la capital de China?
  ID: 7, Score: 30.02, Document: 北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。
  ID: 1, Score: 30.02, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 6, Score: 30.01, Document: Madrid es la capital y la ciudad más grande de España. Situada en el centro de la península ibérica, es el principal centro político, administrativo, económico y cultural del país.
--------------------
クエリ: 韓国的首都是哪座城市？
  ID: 7, Score: 29.51, Document: 北京是中国的首都，也是国家的政治、文化和国际交流中心，位于华北平原的北部。
  ID: 2, Score: 29.40, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。
  ID: 8, Score: 29.27, Document: 서울은 대한민국의 수도이자 최대 도시로, 정치·경제·문화의 중심지 역할을 하고 있습니다.
--------------------
クエリ: 프랑스의 수도는 어디예요?
  ID: 1, Score: 29.27, Document: Paris est la capitale et la plus grande ville de la France. Située sur la Seine, dans le nord du pays, elle est le centre politique, économique et culturel du pays.
  ID: 8, Score: 29.25, Document: 서울은 대한민국의 수도이자 최대 도시로, 정치·경제·문화의 중심지 역할을 하고 있습니다.
  ID: 2, Score: 29.19, Document: 東京（正式には東京都）は、日本の首都である。世界有数の人口を抱える大都市圏であり、日本の行政・金融・商業の中枢としての役割を果たしている。

from pylate import indexes, models, retrieve

model = models.ColBERT(
    model_name_or_path="LiquidAI/LFM2-ColBERT-350M",
)
model.tokenizer.pad_token = model.tokenizer.eos_token

index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,
)

documents_ids = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12"]
documents = [
    "木の温もりあふれるブックカフェで、自家焙煎の深煎りコーヒーと季節のタルトを味わいながら、窓辺から路面電車をのんびり眺められるんだ。",
    "庭にハーブが茂るガーデンカフェでは、ハンドドリップの浅煎りとフレッシュハーブティーが選べて、小鳥のさえずりが BGM 代わりになるよ。",
    "港直送の鯖を炙りしめ鯖にしてくれる専門店、皮目の香ばしさと酢のきりりとした酸味が口いっぱいに広がるんだ。",
    "カウンター割烹の金目鯛の煮付けは、甘辛ダレが骨の隅々まで染みていて、白ご飯が思わずおかわり必至だよね。",
    "昔ながらの屋台ラーメンは鶏ガラの澄んだ醤油スープと細ちぢれ麺が相性抜群で、深夜の胃袋にしみるんだ。",
    "真っ白な豚骨スープに焦がしニンニク油をひと垂らしした濃厚ラーメン、替え玉が無料でつい無限ループしてしまうよ。",
    "スリランカ式の混ぜて食べるプレートカレーでは、15種類のスパイスが複雑に重なって食べ進めるほど香りが花開くんだ。",
    "野菜がごろごろ入った欧風ビーフカレーは、赤ワインとバターのコクが効いたシャバっとルウで後を引くよ。",
    "薪窯ナポリピッツァのマルゲリータは、モッツァレラがびよーんと伸びて焼き立てを頬張る瞬間がたまらない。",
    "４種のチーズをのせたクアトロフォルマッジに蜂蜜を垂らすスタイルが人気で、塩気と甘さのコントラストがクセになるんだ。",
    "しゅわっととろけるバスクチーズケーキ専門店、表面の香ばしい焦げと濃厚クリーミーな中身のギャップが病みつきになるよ。",
    "パリパリの薄皮たい焼きは羽根つきで端っこまで香ばしく、黒あんか白あんか毎回真剣に迷っちゃうんだよね。",
]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,  # ドキュメントの場合は False
    show_progress_bar=True,
)

index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

retriever = retrieve.ColBERT(index=index)

queries = [
    "中華そばが食べたいなぁ",
    "ゆっくりお茶したい",
    "今日はスパイシーな感じで"
]

queries_embeddings = model.encode(
    queries,
    batch_size=32,
    is_query=True, # クエリの場合は False
    show_progress_bar=True,
)

scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=5,
)

for i, query_scores in enumerate(scores):
    print(f"クエリ: {queries[i]}")
    for score_info in query_scores:
        doc_id = score_info['id']
        score = score_info['score']
        doc_index = documents_ids.index(doc_id)
        document_text = documents[doc_index]
        print(f"  ID: {doc_id}, Score: {score:.2f}, Document: {document_text}")
    print("-" * 20)

クエリ: 中華そばが食べたいなぁ
  ID: 5, Score: 28.94, Document: 昔ながらの屋台ラーメンは鶏ガラの澄んだ醤油スープと細ちぢれ麺が相性抜群で、深夜の胃袋にしみるんだ。
  ID: 3, Score: 28.90, Document: 港直送の鯖を炙りしめ鯖にしてくれる専門店、皮目の香ばしさと酢のきりりとした酸味が口いっぱいに広がるんだ。
  ID: 4, Score: 28.89, Document: カウンター割烹の金目鯛の煮付けは、甘辛ダレが骨の隅々まで染みていて、白ご飯が思わずおかわり必至だよね。
  ID: 6, Score: 28.88, Document: 真っ白な豚骨スープに焦がしニンニク油をひと垂らしした濃厚ラーメン、替え玉が無料でつい無限ループしてしまうよ。
  ID: 7, Score: 28.86, Document: スリランカ式の混ぜて食べるプレートカレーでは、15種類のスパイスが複雑に重なって食べ進めるほど香りが花開くんだ。
--------------------
クエリ: ゆっくりお茶したい
  ID: 1, Score: 29.46, Document: 木の温もりあふれるブックカフェで、自家焙煎の深煎りコーヒーと季節のタルトを味わいながら、窓辺から路面電車をのんびり眺められるんだ。
  ID: 2, Score: 29.45, Document: 庭にハーブが茂るガーデンカフェでは、ハンドドリップの浅煎りとフレッシュハーブティーが選べて、小鳥のさえずりが BGM 代わりになるよ。
  ID: 11, Score: 29.24, Document: しゅわっととろけるバスクチーズケーキ専門店、表面の香ばしい焦げと濃厚クリーミーな中身のギャップが病みつきになるよ。
  ID: 3, Score: 29.17, Document: 港直送の鯖を炙りしめ鯖にしてくれる専門店、皮目の香ばしさと酢のきりりとした酸味が口いっぱいに広がるんだ。
  ID: 4, Score: 29.14, Document: カウンター割烹の金目鯛の煮付けは、甘辛ダレが骨の隅々まで染みていて、白ご飯が思わずおかわり必至だよね。
--------------------
クエリ: 今日はスパイシーな感じで
  ID: 7, Score: 29.22, Document: スリランカ式の混ぜて食べるプレートカレーでは、15種類のスパイスが複雑に重なって食べ進めるほど香りが花開くんだ。
  ID: 10, Score: 29.09, Document: ４種のチーズをのせたクアトロフォルマッジに蜂蜜を垂らすスタイルが人気で、塩気と甘さのコントラストがクセになるんだ。
  ID: 3, Score: 29.08, Document: 港直送の鯖を炙りしめ鯖にしてくれる専門店、皮目の香ばしさと酢のきりりとした酸味が口いっぱいに広がるんだ。
  ID: 5, Score: 29.06, Document: 昔ながらの屋台ラーメンは鶏ガラの澄んだ醤油スープと細ちぢれ麺が相性抜群で、深夜の胃袋にしみるんだ。
  ID: 8, Score: 29.05, Document: 野菜がごろごろ入った欧風ビーフカレーは、赤ワインとバターのコクが効いたシャバっとルウで後を引くよ。
--------------------