跨标题、摘要和文本块的多重表示检索

pip install qdrant-client datasets

from datasets import load_dataset

ML_CATEGORIES = {"cs.LG", "cs.CV", "cs.CL", "cs.AI", "stat.ML"}

# Non-streaming so HF caches the parquet locally; first run downloads ~2.5 GB, re-runs are instant.
dataset = load_dataset("gfissore/arxiv-abstracts-2021", split="train")

papers = []
# IDs are roughly chronological; iterate from the end to land on 2021/2020/2019 papers first.
for i in range(len(dataset) - 1, -1, -1):
    if len(papers) >= 20000:
        break
    row = dataset[i]
    if not row["abstract"] or not row["title"]:
        continue
    # categories arrive as space-joined strings (e.g. ["cs.LG cs.CV"]); split each entry.
    cats = [tok for entry in row["categories"] for tok in entry.split()]
    if not any(c in ML_CATEGORIES for c in cats):
        continue
    # Year lives in the YYMM prefix of new-format arXiv IDs ("2104.01234" -> 2021).
    arxiv_id = row["id"]
    if "/" in arxiv_id or "." not in arxiv_id:
        continue  # skip pre-2007 IDs like "math/0506001"
    if 2000 + int(arxiv_id[:2]) < 2018:
        continue
    papers.append({
        "arxiv_id": arxiv_id,
        "title": row["title"].strip(),
        "abstract": row["abstract"].strip(),
        "tags": cats,
    })

from qdrant_client import QdrantClient, models

# Replace url and api_key with your own from https://cloud.qdrant.io
client = QdrantClient(
    url="https://xyz-example.qdrant.io:6333",
    api_key="<your-api-key>",
    cloud_inference=True,
)

# 384 is the output dimension of sentence-transformers/all-minilm-l6-v2, used below for every dense vector.
client.create_collection(
    collection_name="arxiv_multi_repr",
    vectors_config={
        "dense_chunk":   models.VectorParams(size=384, distance=models.Distance.COSINE),
        "dense_title":   models.VectorParams(size=384, distance=models.Distance.COSINE),
        "dense_abstract": models.VectorParams(size=384, distance=models.Distance.COSINE),
    },
    sparse_vectors_config={
        "sparse_title": models.SparseVectorParams(
            modifier=models.Modifier.IDF,
        ),
    },
)

# Index 'document_id' so the Query API can group by it; index 'tags' so we can filter on category.
client.create_payload_index(
    collection_name="arxiv_multi_repr",
    field_name="document_id",
    field_schema=models.PayloadSchemaType.KEYWORD,
)
client.create_payload_index(
    collection_name="arxiv_multi_repr",
    field_name="tags",
    field_schema=models.PayloadSchemaType.KEYWORD,
)

DENSE_MODEL = "sentence-transformers/all-minilm-l6-v2"
BM25_MODEL = "qdrant/bm25"

def chunk_sentences(text, target_len=2):
    """Split text into ~2-sentence chunks; fall back to the full text if it doesn't split cleanly."""
    sentences = [s.strip() for s in text.split(". ") if s.strip()]
    return [". ".join(sentences[i:i + target_len])
            for i in range(0, len(sentences), target_len)] or [text]

points = []
for paper in papers:
    chunks = chunk_sentences(paper["abstract"])

    # Title, abstract, and sparse docs are reused across every chunk of this paper; only the chunk text varies.
    # Cloud Inference embeds each Document on the server, so you don't need a client-side embedding library.
    title_doc   = models.Document(text=paper["title"],    model=DENSE_MODEL)
    abstract_doc = models.Document(text=paper["abstract"], model=DENSE_MODEL)
    # avg_len is the average word count of the indexed text.
    # Default is 256 (document-length); setting it to the actual field length (~10 here) improves BM25 scoring accuracy.
    sparse_doc  = models.Document(
        text=paper["title"],
        model=BM25_MODEL,
        options={"avg_len": 10.0},
    )

    for i, chunk in enumerate(chunks):
        points.append(models.PointStruct(
            id=len(points),
            vector={
                "dense_chunk":     models.Document(text=chunk, model=DENSE_MODEL),
                "dense_title":     title_doc,
                "dense_abstract":   abstract_doc,
                "sparse_title": sparse_doc,
            },
            payload={
                "document_id": paper["arxiv_id"],
                "title":       paper["title"],
                "tags":        paper["tags"],
                "chunk_index": i,
                "chunk_text":  chunk,
            },
        ))

client.upload_points(collection_name="arxiv_multi_repr", points=points, batch_size=256, parallel=2)

def retrieve(query, limit=10, group_size=3, tags=None):
    dense_query  = models.Document(text=query, model=DENSE_MODEL)
    sparse_query = models.Document(text=query, model=BM25_MODEL)
    # Optional category filter. When tags is provided, Qdrant pre-filters candidates
    # to points whose 'tags' payload includes any of the given values.
    query_filter = (
        models.Filter(must=[models.FieldCondition(key="tags", match=models.MatchAny(any=tags))])
        if tags else None
    )
    # query_points_groups runs the prefetches, fuses with RRF, applies the filter, and groups results by document_id.
    # You may need to adjust the per-prefetch limit based on the number of chunks per document; grouping only sees what the prefetch returns.
    return client.query_points_groups(
        collection_name="arxiv_multi_repr",
        prefetch=[
            models.Prefetch(query=dense_query,  using="dense_chunk",    limit=100),
            models.Prefetch(query=dense_query,  using="dense_title",    limit=100),
            models.Prefetch(query=dense_query,  using="dense_abstract", limit=100),
            models.Prefetch(query=sparse_query, using="sparse_title",   limit=100),
        ],
        query=models.FusionQuery(fusion=models.Fusion.RRF),
        query_filter=query_filter,
        group_by="document_id",
        group_size=group_size,
        limit=limit,
    ).groups

query=models.FormulaQuery(
    formula=models.SumExpression(sum=[
        models.MultExpression(mult=[1.0, "$score[0]"]),
        models.MultExpression(mult=[0.5, "$score[1]"]),
        models.MultExpression(mult=[0.4, "$score[2]"]),
        models.MultExpression(mult=[0.3, "$score[3]"]),
    ]),
    defaults={"$score[1]": 0.0, "$score[2]": 0.0, "$score[3]": 0.0},
),