البحث القمعي باستخدام تضمينات ماتريوشكا

عند إنشاء أنظمة بحث متجهات فعالة، يتمثل أحد التحديات الرئيسية في إدارة تكاليف التخزين مع الحفاظ على زمن استجابة واستدعاء مقبول. تُخرج نماذج التضمين الحديثة متجهات ذات مئات أو آلاف الأبعاد، مما يؤدي إلى إنشاء مساحة تخزين كبيرة ونفقات حسابية كبيرة للمتجه الخام والفهرس.

وتقليديًا، يتم تقليل متطلبات التخزين من خلال تطبيق طريقة التكميم أو تقليل الأبعاد قبل إنشاء الفهرس مباشرةً. على سبيل المثال، يمكننا توفير مساحة التخزين عن طريق تقليل الدقة باستخدام التكميم الكمي للمنتج (PQ) أو عدد الأبعاد باستخدام تحليل المكونات الرئيسية (PCA). تقوم هذه الطرق بتحليل مجموعة المتجهات بأكملها للعثور على مجموعة أكثر إحكامًا تحافظ على العلاقات الدلالية بين المتجهات.

رغم فعاليتها، فإن هذه الأساليب القياسية تقلل من الدقة أو الأبعاد مرة واحدة فقط وعلى مقياس واحد. ولكن ماذا لو تمكنا من الحفاظ على طبقات متعددة من التفاصيل في وقت واحد، مثل هرم من التمثيلات المتزايدة الدقة؟

أدخل تضمينات ماتريوشكا. سُميت على اسم الدمى الروسية المتداخلة (انظر الرسم التوضيحي)، هذه التركيبات الذكية تدمج مقاييس متعددة للتمثيل داخل متجه واحد. على عكس طرق المعالجة اللاحقة التقليدية، تتعلم تضمينات ماتريوشكا هذه البنية متعددة المقاييس أثناء عملية التدريب الأولية. والنتيجة ملحوظة: لا يقتصر الأمر على أن التضمين الكامل لا يلتقط دلالات المدخلات فحسب، بل إن كل بادئة مجموعة فرعية متداخلة (النصف الأول، الربع الأول، إلخ) توفر تمثيلاً متماسكاً، وإن كان أقل تفصيلاً.

في هذا الدفتر، ندرس كيفية استخدام تضمينات ماتريوشكا مع ميلفوس للبحث الدلالي. نوضح خوارزمية تسمى "البحث القمعي" التي تسمح لنا بإجراء بحث التشابه على مجموعة فرعية صغيرة من أبعاد التضمين دون انخفاض كبير في التذكر.

import functools

from datasets import load_dataset
import numpy as np
import pandas as pd
import pymilvus
from pymilvus import MilvusClient
from pymilvus import FieldSchema, CollectionSchema, DataType
from sentence_transformers import SentenceTransformer
import torch
import torch.nn.functional as F
from tqdm import tqdm

تحميل نموذج تضمين ماتريوشكا

بدلاً من استخدام نموذج تضمين قياسي مثل sentence-transformers/all-MiniLM-L12-v2، نستخدم نموذجًا من Nomic مدربًا خصيصًا لإنتاج تضمينات ماتريوشكا.

model = SentenceTransformer(
    # Remove 'device='mps' if running on non-Mac device
    "nomic-ai/nomic-embed-text-v1.5",
    trust_remote_code=True,
    device="mps",
)

<All keys matched successfully>

تحميل مجموعة البيانات، وتضمين العناصر، وبناء قاعدة بيانات المتجهات

الكود التالي هو تعديل لذلك من صفحة التوثيق "البحث عن الأفلام باستخدام محولات الجمل وميلفوس". أولاً، نقوم بتحميل مجموعة البيانات من HuggingFace. وهي تحتوي على حوالي 35 ألف مدخل، كل منها يتوافق مع فيلم له مقالة في ويكيبيديا. سنستخدم الحقلين Title و PlotSummary في هذا المثال.

ds = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries", split="train")
print(ds)

Dataset({
    features: ['Release Year', 'Title', 'Origin/Ethnicity', 'Director', 'Cast', 'Genre', 'Wiki Page', 'Plot', 'PlotSummary'],
    num_rows: 34886
})

بعد ذلك، نتصل بقاعدة بيانات Milvus Lite، ونحدد مخطط البيانات، وننشئ مجموعة بهذا المخطط. سنقوم بتخزين كل من التضمين غير الطبيعي والسدس الأول من التضمين في حقول منفصلة. والسبب في ذلك هو أننا نحتاج إلى أول 1/6 من تضمين ماتريوشكا لإجراء بحث التشابه، والخمسة أسداس المتبقية من التضمينات لإعادة ترتيب نتائج البحث وتحسينها.

embedding_dim = 768
search_dim = 128
collection_name = "movie_embeddings"

client = MilvusClient(uri="./wiki-movie-plots-matryoshka.db")

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
    # First sixth of unnormalized embedding vector
    FieldSchema(name="head_embedding", dtype=DataType.FLOAT_VECTOR, dim=search_dim),
    # Entire unnormalized embedding vector
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
]

schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
client.create_collection(collection_name=collection_name, schema=schema)

لا يدعم برنامج Milvus حاليًا البحث على مجموعات فرعية من التضمينات، لذلك نقوم بتقسيم التضمينات إلى جزأين: يمثل الرأس المجموعة الفرعية الأولية من المتجه للفهرسة والبحث، والذيل هو الباقي. تم تدريب النموذج على البحث عن تشابه مسافة جيب التمام، لذلك نقوم بتطبيع تضمينات الرأس. ومع ذلك، من أجل حساب أوجه التشابه لمجموعات فرعية أكبر لاحقًا، نحتاج إلى تخزين معيار تضمين الرأس، حتى نتمكن من عدم تطبيعه قبل الانضمام إلى الذيل.

لإجراء البحث عن طريق أول 1/6 من التضمين الأول، سنحتاج إلى إنشاء فهرس بحث متجه على الحقل head_embedding. في وقت لاحق، سنقارن نتائج "البحث القمعي" بالبحث المتجه العادي، ومن ثم ننشئ فهرس بحث على التضمين الكامل أيضًا.

من المهم أن نستخدم COSINE بدلاً من مقياس المسافة IP ، لأنه بخلاف ذلك سنحتاج إلى تتبع معايير التضمين، مما سيعقد التنفيذ (سيكون هذا أكثر منطقية بمجرد وصف خوارزمية البحث القمعي).

index_params = client.prepare_index_params()
index_params.add_index(
    field_name="head_embedding", index_type="FLAT", metric_type="COSINE"
)
index_params.add_index(field_name="embedding", index_type="FLAT", metric_type="COSINE")
client.create_index(collection_name, index_params)

أخيراً، نقوم بترميز ملخصات الحبكة لجميع الأفلام الـ 35 ألفاً وإدخال التضمينات المقابلة في قاعدة البيانات.

for batch in tqdm(ds.batch(batch_size=512)):
    # This particular model requires us to prefix 'search_document:' to stored entities
    plot_summary = ["search_document: " + x.strip() for x in batch["PlotSummary"]]

    # Output of embedding model is unnormalized
    embeddings = model.encode(plot_summary, convert_to_tensor=True)
    head_embeddings = embeddings[:, :search_dim]

    data = [
        {
            "title": title,
            "head_embedding": head.cpu().numpy(),
            "embedding": embedding.cpu().numpy(),
        }
        for title, head, embedding in zip(batch["Title"], head_embeddings, embeddings)
    ]
    res = client.insert(collection_name=collection_name, data=data)

100%|██████████| 69/69 [05:57<00:00,  5.18s/it]

إجراء بحث قمعي

لننفذ الآن "بحثاً قمعياً" باستخدام أول 1/6 من أبعاد تضمين ماتريوشكا. لديّ ثلاثة أفلام في الاعتبار لاسترجاعها وقمت بإنتاج ملخص الرسم البياني الخاص بي للاستعلام عن قاعدة البيانات. نقوم بتضمين الاستعلامات، ثم نجري بحثًا متجهًا على الحقل head_embedding ، ونسترجع 128 نتيجة مرشحة.

queries = [
    "An archaeologist searches for ancient artifacts while fighting Nazis.",
    "A teenager fakes illness to get off school and have adventures with two friends.",
    "A young couple with a kid look after a hotel during winter and the husband goes insane.",
]


# Search the database based on input text
def embed_search(data):
    embeds = model.encode(data)
    return [x for x in embeds]


# This particular model requires us to prefix 'search_query:' to queries
instruct_queries = ["search_query: " + q.strip() for q in queries]
search_data = embed_search(instruct_queries)

# Normalize head embeddings
head_search = [x[:search_dim] for x in search_data]

# Perform standard vector search on first sixth of embedding dimensions
res = client.search(
    collection_name=collection_name,
    data=head_search,
    anns_field="head_embedding",
    limit=128,
    output_fields=["title", "head_embedding", "embedding"],
)

عند هذه النقطة، نكون قد أجرينا بحثًا على مساحة متجه أصغر بكثير، وبالتالي من المحتمل أن يكون زمن الاستجابة أقل ومتطلبات تخزين أقل للفهرس مقارنة بالبحث على المساحة الكاملة. دعونا نفحص أفضل 5 مطابقات لكل استعلام:

for query, hits in zip(queries, res):
    rows = [x["entity"] for x in hits][:5]

    print("Query:", query)
    print("Results:")
    for row in rows:
        print(row["title"].strip())
    print()

Query: An archaeologist searches for ancient artifacts while fighting Nazis.
Results:
"Pimpernel" Smith
Black Hunters
The Passage
Counterblast
Dominion: Prequel to the Exorcist

Query: A teenager fakes illness to get off school and have adventures with two friends.
Results:
How to Deal
Shorts
Blackbird
Valentine
Unfriended

Query: A young couple with a kid look after a hotel during winter and the husband goes insane.
Results:
Ghostkeeper
Our Vines Have Tender Grapes
The Ref
Impact
The House in Marsh Road

كما نرى، تأثر التذكر نتيجة لاقتطاع التضمينات أثناء البحث. يعمل البحث القمعي على إصلاح ذلك باستخدام حيلة ذكية: يمكننا استخدام ما تبقى من أبعاد التضمين لإعادة ترتيب قائمة المرشحين وتشذيبها لاستعادة أداء الاسترجاع دون إجراء أي عمليات بحث إضافية مكلفة عن المتجهات.

لتسهيل شرح خوارزمية البحث القمعي نقوم بتحويل نتائج بحث ميلفوس لكل استعلام إلى إطار بيانات بانداس.

def hits_to_dataframe(hits: pymilvus.client.abstract.Hits) -> pd.DataFrame:
    """
    Convert a Milvus search result to a Pandas dataframe. This function is specific to our data schema.

    """
    rows = [x["entity"] for x in hits]
    rows_dict = [
        {"title": x["title"], "embedding": torch.tensor(x["embedding"])} for x in rows
    ]
    return pd.DataFrame.from_records(rows_dict)


dfs = [hits_to_dataframe(hits) for hits in res]

والآن، لإجراء البحث القمعي نقوم بالتكرار على مجموعات فرعية أكبر بشكل متزايد من التضمينات. في كل عملية تكرار، نعيد ترتيب المرشحين وفقًا لأوجه التشابه الجديدة ونشذب جزءًا من أقلها مرتبة.

لتوضيح ذلك، استرجعنا من الخطوة السابقة 128 مرشحًا باستخدام 1/6 من أبعاد التضمين والاستعلام. تتمثل الخطوة الأولى في إجراء البحث القمعي في إعادة حساب أوجه التشابه بين الاستعلامات والمرشحين باستخدام الثلث الأول من الأبعاد. يتم تشذيب ال 64 مرشحًا السفلي. ثم نكرر هذه العملية مع أول 2/3 من الأبعاد، ثم جميع الأبعاد، مع تشذيب 32 و16 مرشحًا على التوالي.

# An optimized implementation would vectorize the calculation of similarity scores across rows (using a matrix)
def calculate_score(row, query_emb=None, dims=768):
    emb = F.normalize(row["embedding"][:dims], dim=-1)
    return (emb @ query_emb).item()


# You could also add a top-K parameter as a termination condition
def funnel_search(
    df: pd.DataFrame, query_emb, scales=[256, 512, 768], prune_ratio=0.5
) -> pd.DataFrame:
    # Loop over increasing prefixes of the embeddings
    for dims in scales:
        # Query vector must be normalized for each new dimensionality
        emb = torch.tensor(query_emb[:dims] / np.linalg.norm(query_emb[:dims]))

        # Score
        scores = df.apply(
            functools.partial(calculate_score, query_emb=emb, dims=dims), axis=1
        )
        df["scores"] = scores

        # Re-rank
        df = df.sort_values(by="scores", ascending=False)

        # Prune (in our case, remove half of candidates at each step)
        df = df.head(int(prune_ratio * len(df)))

    return df


dfs_results = [
    {"query": query, "results": funnel_search(df, query_emb)}
    for query, df, query_emb in zip(queries, dfs, search_data)
]

for d in dfs_results:
    print(d["query"], "\n", d["results"][:5]["title"], "\n")

An archaeologist searches for ancient artifacts while fighting Nazis. 
 0           "Pimpernel" Smith
1               Black Hunters
29    Raiders of the Lost Ark
34             The Master Key
51            My Gun Is Quick
Name: title, dtype: object 

A teenager fakes illness to get off school and have adventures with two friends. 
 21               How I Live Now
32     On the Edge of Innocence
77             Bratz: The Movie
4                    Unfriended
108                  Simon Says
Name: title, dtype: object 

A young couple with a kid look after a hotel during winter and the husband goes insane. 
 9         The Shining
0         Ghostkeeper
11     Fast and Loose
7      Killing Ground
12         Home Alone
Name: title, dtype: object 

لقد تمكنا من استعادة التذكر دون إجراء أي عمليات بحث إضافية للمتجهات! من الناحية النوعية، يبدو أن هذه النتائج تتمتع باستدعاء أعلى لفيلمي "Raiders of the Lost Ark" و"The Shining" من البحث المتجه القياسي في البرنامج التعليمي، "البحث عن الأفلام باستخدام ميلفوس ومحوّلات الجمل"، والذي يستخدم نموذج تضمين مختلف. ومع ذلك، فإنه غير قادر على العثور على فيلم "Ferris Bueller's Day Off"، والذي سنعود إليه لاحقًا في هذا الدفتر. (راجع بحث "تعلم تمثيل ماتريوشكا " للاطلاع على المزيد من التجارب الكمية والقياس المعياري).

مقارنة البحث القمعي بالبحث العادي

لنقارن نتائج بحثنا القمعي بالبحث المتجه القياسي على نفس مجموعة البيانات بنفس نموذج التضمين. نجري بحثًا على التضمينات الكاملة.

# Search on entire embeddings
res = client.search(
    collection_name=collection_name,
    data=search_data,
    anns_field="embedding",
    limit=5,
    output_fields=["title", "embedding"],
)

for query, hits in zip(queries, res):
    rows = [x["entity"] for x in hits]

    print("Query:", query)
    print("Results:")
    for row in rows:
        print(row["title"].strip())
    print()

Query: An archaeologist searches for ancient artifacts while fighting Nazis.
Results:
"Pimpernel" Smith
Black Hunters
Raiders of the Lost Ark
The Master Key
My Gun Is Quick

Query: A teenager fakes illness to get off school and have adventures with two friends.
Results:
A Walk to Remember
Ferris Bueller's Day Off
How I Live Now
On the Edge of Innocence
Bratz: The Movie

Query: A young couple with a kid look after a hotel during winter and the husband goes insane.
Results:
The Shining
Ghostkeeper
Fast and Loose
Killing Ground
Home Alone

باستثناء نتائج البحث عن "مراهق يتظاهر بالمرض للتملص من المدرسة..."، فإن النتائج في إطار البحث القمعي مطابقة تقريبًا للبحث الكامل، على الرغم من أن البحث القمعي أُجري على مساحة بحث من 128 بُعدًا مقابل 768 بُعدًا للبحث العادي.

التحقيق في فشل استرجاع البحث القمعي في يوم عطلة فيريس بويلر

لماذا لم ينجح البحث القمعي في استرجاع فيلم Ferris Bueller's Day Off؟ دعونا نفحص ما إذا كان موجودًا في قائمة المرشحين الأصلية أم أنه تمت تصفيته عن طريق الخطأ.

queries2 = [
    "A teenager fakes illness to get off school and have adventures with two friends."
]


# Search the database based on input text
def embed_search(data):
    embeds = model.encode(data)
    return [x for x in embeds]


instruct_queries = ["search_query: " + q.strip() for q in queries2]
search_data2 = embed_search(instruct_queries)
head_search2 = [x[:search_dim] for x in search_data2]

# Perform standard vector search on subset of embeddings
res = client.search(
    collection_name=collection_name,
    data=head_search2,
    anns_field="head_embedding",
    limit=256,
    output_fields=["title", "head_embedding", "embedding"],
)

for query, hits in zip(queries, res):
    rows = [x["entity"] for x in hits]

    print("Query:", queries2[0])
    for idx, row in enumerate(rows):
        if row["title"].strip() == "Ferris Bueller's Day Off":
            print(f"Row {idx}: Ferris Bueller's Day Off")

Query: A teenager fakes illness to get off school and have adventures with two friends.
Row 228: Ferris Bueller's Day Off

نرى أن المشكلة تكمن في أن قائمة المرشحين الأولية لم تكن كبيرة بما فيه الكفاية، أو بالأحرى أن النتيجة المطلوبة ليست مشابهة بما فيه الكفاية للاستعلام على أعلى مستوى من التفصيل. يؤدي تغييره من 128 إلى 256 إلى استرجاع ناجح. يجب أن نشكل قاعدة عامة لتعيين عدد المرشحين في مجموعة المرشحين المحتجزة لتقييم المفاضلة بين الاسترجاع وزمن الاستجابة تجريبيًا.

dfs = [hits_to_dataframe(hits) for hits in res]

dfs_results = [
    {"query": query, "results": funnel_search(df, query_emb)}
    for query, df, query_emb in zip(queries2, dfs, search_data2)
]

for d in dfs_results:
    print(d["query"], "\n", d["results"][:7]["title"].to_string(index=False), "\n")

A teenager fakes illness to get off school and have adventures with two friends. 
       A Walk to Remember
Ferris Bueller's Day Off
          How I Live Now
On the Edge of Innocence
        Bratz: The Movie
              Unfriended
              Simon Says 

هل الترتيب مهم؟ تضمين البادئة مقابل تضمين اللاحقة.

لقد تم تدريب النموذج على الأداء الجيد لمطابقة البادئات الأصغر من التضمينات بشكل متكرر. هل ترتيب الأبعاد التي نستخدمها مهم؟ على سبيل المثال، هل يمكننا أيضًا أخذ مجموعات فرعية من التضمينات التي هي عبارة عن لاحقات؟ في هذه التجربة، نقوم بعكس ترتيب الأبعاد في تضمينات ماتريوشكا ونجري بحثًا قمعيًا.

client = MilvusClient(uri="./wikiplots-matryoshka-flipped.db")

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
    FieldSchema(name="head_embedding", dtype=DataType.FLOAT_VECTOR, dim=search_dim),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
]

schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
client.create_collection(collection_name=collection_name, schema=schema)

index_params = client.prepare_index_params()
index_params.add_index(
    field_name="head_embedding", index_type="FLAT", metric_type="COSINE"
)
client.create_index(collection_name, index_params)

huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
    - Avoid using `tokenizers` before the fork if possible
    - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)

for batch in tqdm(ds.batch(batch_size=512)):
    plot_summary = ["search_document: " + x.strip() for x in batch["PlotSummary"]]

    # Encode and flip embeddings
    embeddings = model.encode(plot_summary, convert_to_tensor=True)
    embeddings = torch.flip(embeddings, dims=[-1])
    head_embeddings = embeddings[:, :search_dim]

    data = [
        {
            "title": title,
            "head_embedding": head.cpu().numpy(),
            "embedding": embedding.cpu().numpy(),
        }
        for title, head, embedding in zip(batch["Title"], head_embeddings, embeddings)
    ]
    res = client.insert(collection_name=collection_name, data=data)

100%|██████████| 69/69 [05:50<00:00,  5.08s/it]

# Normalize head embeddings

flip_search_data = [
    torch.flip(torch.tensor(x), dims=[-1]).cpu().numpy() for x in search_data
]
flip_head_search = [x[:search_dim] for x in flip_search_data]

# Perform standard vector search on subset of embeddings
res = client.search(
    collection_name=collection_name,
    data=flip_head_search,
    anns_field="head_embedding",
    limit=128,
    output_fields=["title", "head_embedding", "embedding"],
)

dfs = [hits_to_dataframe(hits) for hits in res]

dfs_results = [
    {"query": query, "results": funnel_search(df, query_emb)}
    for query, df, query_emb in zip(queries, dfs, flip_search_data)
]

for d in dfs_results:
    print(
        d["query"],
        "\n",
        d["results"][:7]["title"].to_string(index=False, header=False),
        "\n",
    )

An archaeologist searches for ancient artifacts while fighting Nazis. 
       "Pimpernel" Smith
          Black Hunters
Raiders of the Lost Ark
         The Master Key
        My Gun Is Quick
            The Passage
        The Mole People 

A teenager fakes illness to get off school and have adventures with two friends. 
                       A Walk to Remember
                          How I Live Now
                              Unfriended
Cirque du Freak: The Vampire's Assistant
                             Last Summer
                                 Contest
                                 Day One 

A young couple with a kid look after a hotel during winter and the husband goes insane. 
         Ghostkeeper
     Killing Ground
Leopard in the Snow
              Stone
          Afterglow
         Unfaithful
     Always a Bride 

يكون التذكر أضعف بكثير من البحث القمعي أو البحث العادي كما هو متوقع (تم تدريب نموذج التضمين عن طريق التعلم التبايني على بادئات أبعاد التضمين وليس اللواحق).

ملخص

فيما يلي مقارنة لنتائج بحثنا عبر الطرق:

لقد أوضحنا كيفية استخدام تضمينات ماتريوشكا مع ميلفوس لإجراء خوارزمية بحث دلالي أكثر كفاءة تسمى "البحث القمعي". كما استكشفنا أيضًا أهمية خطوات إعادة الترتيب والتشذيب في الخوارزمية، بالإضافة إلى وضع الفشل عندما تكون قائمة المرشحين الأولية صغيرة جدًا. وأخيرًا، ناقشنا مدى أهمية ترتيب الأبعاد عند تكوين التضمينات الفرعية - يجب أن يكون بنفس الطريقة التي تم تدريب النموذج عليها. أو بالأحرى، فقط لأن النموذج قد تم تدريبه بطريقة معينة بحيث تكون البادئات من التضمينات ذات معنى. أنت تعرف الآن كيفية تنفيذ تضمينات ماتريوشكا والبحث القمعي لتقليل تكاليف تخزين البحث الدلالي دون التضحية بالكثير من أداء الاسترجاع!