search/bm25.py at main · abhishekkrthakur/search · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
from vespa.package import (
    ApplicationPackage,
    Document,
    Field,
    FieldSet,
    Function,
    RankProfile,
    Schema,
)
from vespa.deployment import VespaDocker
from vespa.io import VespaResponse
from datasets import load_dataset
from tqdm import tqdm

package = ApplicationPackage(
    name="simplesearch",
    schema=[
        Schema(
            name="doc",
            document=Document(
                fields=[
                    Field(
                        name="id",
                        type="string",
                        indexing=["summary"],
                    ),
                    Field(
                        name="text",
                        type="string",
                        indexing=["index", "summary"],
                        index="enable-bm25",
                    ),
                    Field(
                        name="url",
                        type="string",
                        indexing=["index", "summary"],
                        index="enable-bm25",
                    ),
                ]
            ),
            fieldsets=[
                FieldSet(name="default", fields=["text", "url"]),
            ],
            rank_profiles=[
                # 1) BM25 only on text
                RankProfile(
                    name="bm25_text_only",
                    functions=[
                        Function(
                            name="bm25text",
                            expression="bm25(text)",
                        ),
                    ],
                    first_phase="bm25text",
                ),
                # 2) BM25 only on url
                RankProfile(
                    name="bm25_url_only",
                    functions=[
                        Function(
                            name="bm25url",
                            expression="bm25(url)",
                        ),
                    ],
                    first_phase="bm25url",
                ),
                # 3) Original combined BM25 (defaults for k and b)
                RankProfile(
                    name="bm25",
                    functions=[
                        Function(
                            name="bm25texturl",
                            expression="bm25(text) + 0.1 * bm25(url)",
                        ),
                    ],
                    first_phase="bm25texturl",
                ),
                # --- 4) Combined BM25 with different k and b ---
                RankProfile(
                    name="bm25_comb_tuned",
                    functions=[
                        Function(
                            name="bm25texturl_tuned",
                            expression="bm25(text) + 0.1 * bm25(url)",
                        ),
                    ],
                    first_phase="bm25texturl_tuned",
                    rank_properties=[
                        ("bm25(text).k1", "1.8"),
                        ("bm25(text).b", "0.40"),
                        ("bm25(url).k1", "0.9"),
                        ("bm25(url).b", "0.30"),
                    ],
                ),
            ],
        ),
    ],
)