90 lines
2.6 KiB
Go
90 lines
2.6 KiB
Go
package digest
|
||
|
||
// P5b task 01: BM25, the lexical ranking used while a long document's embeddings are still
|
||
// being built (and inside the hybrid score after).
|
||
|
||
import (
|
||
"math"
|
||
"strings"
|
||
"testing"
|
||
)
|
||
|
||
func TestBM25RanksTheRareTerm(t *testing.T) {
|
||
docs := []string{
|
||
"the ship sailed from helsinki to tallinn in the morning",
|
||
"the ship stayed in port all day",
|
||
"a zanzibarite sample was found near the harbour",
|
||
"the harbour master logged the ship",
|
||
}
|
||
s := NewBM25(docs).Scores("where was the zanzibarite found")
|
||
if len(s) != len(docs) {
|
||
t.Fatalf("%d scores for %d docs", len(s), len(docs))
|
||
}
|
||
best := 0
|
||
for i := range s {
|
||
if s[i] > s[best] {
|
||
best = i
|
||
}
|
||
}
|
||
if best != 2 {
|
||
t.Errorf("best = doc %d (%v), want doc 2, the only one with the rare term", best, s)
|
||
}
|
||
}
|
||
|
||
func TestBM25Properties(t *testing.T) {
|
||
filler := strings.Repeat("word ", 50)
|
||
docs := []string{
|
||
"alpha " + filler, // 0: tf 1
|
||
strings.Repeat("alpha ", 10) + filler, // 1: tf 10, same filler
|
||
"alpha common", // 2: tf 1, short
|
||
"common " + filler, // 3: no alpha
|
||
}
|
||
idx := NewBM25(docs)
|
||
s := idx.Scores("alpha")
|
||
if !(s[1] > s[0]) {
|
||
t.Errorf("more occurrences must score higher: tf10 %v vs tf1 %v", s[1], s[0])
|
||
}
|
||
if !(s[1] < 10*s[0]) {
|
||
t.Errorf("term frequency must saturate: tf10 %v, 10 × tf1 %v", s[1], 10*s[0])
|
||
}
|
||
if !(s[2] > s[0]) {
|
||
t.Errorf("the same term in a shorter document must score higher: short %v, long %v", s[2], s[0])
|
||
}
|
||
if s[3] != 0 {
|
||
t.Errorf("a document without the term scores %v, want 0", s[3])
|
||
}
|
||
// IDF: a term in every document is worth less than a rare one.
|
||
all := NewBM25([]string{"xx yy", "xx zz", "xx ww"}).Scores("xx")
|
||
rare := NewBM25([]string{"xx yy", "qq zz", "qq ww"}).Scores("xx")
|
||
if !(rare[0] > all[0]) {
|
||
t.Errorf("idf: a term in 1 of 3 docs must outscore one in 3 of 3: %v vs %v", rare[0], all[0])
|
||
}
|
||
}
|
||
|
||
func TestBM25Degenerate(t *testing.T) {
|
||
for _, docs := range [][]string{nil, {}, {""}, {"", "a b"}} {
|
||
s := NewBM25(docs).Scores("a")
|
||
if len(s) != len(docs) {
|
||
t.Errorf("%q: %d scores", docs, len(s))
|
||
}
|
||
for _, v := range s {
|
||
if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
|
||
t.Errorf("%q: score %v", docs, v)
|
||
}
|
||
}
|
||
}
|
||
for _, v := range NewBM25([]string{"a b", "c d"}).Scores("") {
|
||
if v != 0 {
|
||
t.Errorf("empty query scored %v", v)
|
||
}
|
||
}
|
||
// Deterministic.
|
||
a := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
|
||
b := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
|
||
for i := range a {
|
||
if a[i] != b[i] {
|
||
t.Fatal("scores differ between identical runs")
|
||
}
|
||
}
|
||
}
|