Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement BM25 scoring algorithm #135

Merged
merged 20 commits into from
Feb 23, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 39 additions & 12 deletions DESIGN_DOCUMENT.md
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,9 @@ fact necessarily go against the primary project goals.
`MiniSearch` is composed of two layers:

1. A compact and versatile data structure for indexing terms, providing
lookup by exact match, prefix match, and fuzzy match
lookup by exact match, prefix match, and fuzzy match.
2. An API layer on top of this data structure, providing the search
features
features.

Here follows a description of these two layers.

Expand Down Expand Up @@ -157,17 +157,44 @@ indexed within the same index, to further save space. The index is therefore
structured as following:

```
term -> field -> { document frequency, posting list }
term -> field -> document -> term frequency
lucaong marked this conversation as resolved.
Show resolved Hide resolved
```

The fields and documents are referenced in the index with a short numeric ID for
performance and to save space.

### Search result scoring

When performing a search, the entries corresponding to the search term are
looked up in the index (optionally searching the index with prefix or fuzzy
search), then the documents are scored with a variant of
[Tf-Idf](https://en.wikipedia.org/wiki/Tf–idf), and finally results for
different search terms are merged with the given combinator function (by default
`OR`, but `AND` can be specified).

As the document IDs necessarily occur many times in the posting list, as a space
optimization they are substituted by short generated IDs. An index of short ID
to original ID is maintained alongside the search index, to reconstruct the
original IDs. A similar optimization is applied to the field names.
search). If the combination of term, field and document is found, then this
indicates that the term was present in this particular document field. But it is
not helpful to return all matching documents in an arbitrary order. We want to
return the results in order of _relevance_.

For every document field matching a term, a relevance score is calculated. It
indicates the quality of the match, with a higher score indicating a better
match. The variables that are used to calculate the score are:
- The frequency of the term in the document field that is being scored.
- The total number of documents with matching fields for this term.
- The total number of indexed documents.
- The length of this field.
- The average length of this field for all indexed documents.

The scoring algorithm is based on
[BM25](https://en.wikipedia.org/wiki/Okapi_BM25) (and its derivative BM25+),
which is also used in other popular search engines such as Lucene. BM25 is an
improvement on [TF-IDF](https://en.wikipedia.org/wiki/Tf–idf) and incorporates
the following ideas:
- If a term is less common, the score should be higher (like TD-IDF).
- If a term occurs more frequently, the score should be higher (so far this is
the same as TD-IDF). But the relationship is not linear. If a term occurs
twice as often, the score is _not_ twice as high.
- If a document field is shorter, it requires fewer term occurrences to be
achieve the same relevance as a longer document field. This encodes the idea
that a term occurring once in, say, a title is more relevant than a word
occuring once in a long paragraph.

The scores are calculated for every document field matching a query term. The
results are added. To reward documents that match the most terms, the final
score is multiplied by the number of matching terms in the query.
10 changes: 6 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,15 @@ UIs, removing the need to make requests to a search server.
* Memory-efficient index, designed to support memory-constrained use cases
like mobile browsers.

* Exact match, prefix search, fuzzy match, field boosting
* Exact match, prefix search, fuzzy match, field boosting.

* Auto-suggestion engine, for auto-completion of search queries
* Auto-suggestion engine, for auto-completion of search queries.

* Documents can be added and removed from the index at any time
* Modern search result ranking algorithm.

* Zero external dependencies
* Documents can be added and removed from the index at any time.

* Zero external dependencies.

`MiniSearch` strives to expose a simple API that provides the building blocks to
build custom solutions, while keeping a small and well tested codebase.
Expand Down
225 changes: 221 additions & 4 deletions src/MiniSearch.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ describe('MiniSearch', () => {
expect(ms._fieldIds).toEqual({ title: 0, text: 1 })
expect(ms._documentIds.size).toEqual(0)
expect(ms._fieldLength.size).toEqual(0)
expect(ms._averageFieldLength.length).toEqual(0)
expect(ms._avgFieldLength.length).toEqual(0)
expect(ms._options).toMatchObject(options)
})
})
Expand Down Expand Up @@ -163,14 +163,14 @@ describe('MiniSearch', () => {
it('cleans up all data of the deleted document', () => {
const otherDocument = { id: 4, title: 'Decameron', text: 'Umana cosa è aver compassione degli afflitti' }
const originalFieldLength = new Map(ms._fieldLength)
const originalAverageFieldLength = ms._averageFieldLength.slice()
const originalAverageFieldLength = ms._avgFieldLength.slice()

ms.add(otherDocument)
ms.remove(otherDocument)

expect(ms.documentCount).toEqual(3)
expect(ms._fieldLength).toEqual(originalFieldLength)
expect(ms._averageFieldLength).toEqual(originalAverageFieldLength)
expect(ms._avgFieldLength).toEqual(originalAverageFieldLength)
})

it('does not remove terms from other documents', () => {
Expand Down Expand Up @@ -618,7 +618,7 @@ describe('MiniSearch', () => {
const results = ms.search(query, { boostDocument })
expect(boostDocument).toHaveBeenCalledWith(1, 'divina')
expect(boostDocument).toHaveBeenCalledWith(1, 'commedia')
expect(results[0].score).toEqual(resultsWithoutBoost[0].score * boostFactor)
expect(results[0].score).toBeCloseTo(resultsWithoutBoost[0].score * boostFactor)
})

it('skips document if boostDocument returns a falsy value', () => {
Expand Down Expand Up @@ -797,6 +797,223 @@ describe('MiniSearch', () => {
})
})
})

describe('movie ranking set', () => {
const ms = new MiniSearch({
fields: ['title', 'description'],
storeFields: ['title']
})

ms.add({
id: 'tt1487931',
title: 'Khumba',
description: 'When half-striped zebra Khumba is blamed for the lack of rain by the rest of his insular, superstitious herd, he embarks on a daring quest to earn his stripes. In his search for the legendary waterhole in which the first zebras got their stripes, Khumba meets a quirky range of characters and teams up with an unlikely duo: overprotective wildebeest Mama V and Bradley, a self-obsessed, flamboyant ostrich. But before he can reunite with his herd, Khumba must confront Phango, a sadistic leopard who controls the waterholes and terrorizes all the animals in the Great Karoo. It\'s not all black-and-white in this colorful adventure with a difference.'
})

ms.add({
id: 'tt8737608',
title: 'Rams',
description: 'A feud between two sheep farmers.'
})

ms.add({
id: 'tt0983983',
title: 'Shaun the Sheep',
description: 'Shaun is a cheeky and mischievous sheep at Mossy Bottom farm who\'s the leader of the flock and always plays slapstick jokes, pranks and causes trouble especially on Farmer X and his grumpy guide dog, Bitzer.'
})

ms.add({
id: 'tt5174284',
title: 'Shaun the Sheep: The Farmer\'s Llamas',
description: 'At the annual County Fair, three peculiar llamas catch the eye of Shaun, who tricks the unsuspecting Farmer into buying them. At first, it\'s all fun and games at Mossy Bottom Farm until the trio of unruly animals shows their true colours, wreaking havoc before everyone\'s eyes. Now, it\'s up to Bitzer and Shaun to come up with a winning strategy, if they want to reclaim the farm. Can they rid the once-peaceful ranch of the troublemakers?'
})

ms.add({
id: 'tt0102926',
title: 'The Silence of the Lambs',
description: 'F.B.I. trainee Clarice Starling (Jodie Foster) works hard to advance her career, while trying to hide or put behind her West Virginia roots, of which if some knew, would automatically classify her as being backward or white trash. After graduation, she aspires to work in the agency\'s Behavioral Science Unit under the leadership of Jack Crawford (Scott Glenn). While she is still a trainee, Crawford asks her to question Dr. Hannibal Lecter (Sir Anthony Hopkins), a psychiatrist imprisoned, thus far, for eight years in maximum security isolation for being a serial killer who cannibalized his victims. Clarice is able to figure out the assignment is to pick Lecter\'s brains to help them solve another serial murder case, that of someone coined by the media as "Buffalo Bill" (Ted Levine), who has so far killed five victims, all located in the eastern U.S., all young women, who are slightly overweight (especially around the hips), all who were drowned in natural bodies of water, and all who were stripped of large swaths of skin. She also figures that Crawford chose her, as a woman, to be able to trigger some emotional response from Lecter. After speaking to Lecter for the first time, she realizes that everything with him will be a psychological game, with her often having to read between the very cryptic lines he provides. She has to decide how much she will play along, as his request in return for talking to him is to expose herself emotionally to him. The case takes a more dire turn when a sixth victim is discovered, this one from who they are able to retrieve a key piece of evidence, if Lecter is being forthright as to its meaning. A potential seventh victim is high profile Catherine Martin (Brooke Smith), the daughter of Senator Ruth Martin (Diane Baker), which places greater scrutiny on the case as they search for a hopefully still alive Catherine. Who may factor into what happens is Dr. Frederick Chilton (Anthony Heald), the warden at the prison, an opportunist who sees the higher profile with Catherine, meaning a higher profile for himself if he can insert himself successfully into the proceedings.'
})

ms.add({
id: 'tt0395479',
title: 'Boundin\'',
description: 'In the not too distant past, a lamb lives in the desert plateau just below the snow line. He is proud of how bright and shiny his coat of wool is, so much so that it makes him want to dance, which in turn makes all the other creatures around him also want to dance. His life changes when one spring day he is captured, his wool shorn, and thrown back out onto the plateau all naked and pink. But a bounding jackalope who wanders by makes the lamb look at life a little differently in seeing that there is always something exciting in life to bound about.'
})

ms.add({
id: 'tt9812474',
title: 'Lamb',
description: 'Haunted by the indelible mark of loss and silent grief, sad-eyed María and her taciturn husband, Ingvar, seek solace in back-breaking work and the demanding schedule at their sheep farm in the remote, harsh, wind-swept landscapes of mountainous Iceland. Then, with their relationship hanging on by a thread, something unexplainable happens, and just like that, happiness blesses the couple\'s grim household once more. Now, as a painful ending gives birth to a new beginning, Ingvar\'s troubled brother, Pétur, arrives at the farmhouse, threatening María and Ingvar\'s delicate, newfound bliss. But, nature\'s gifts demand sacrifice. How far are ecstatic María and Ingvar willing to go in the name of love?'
})

ms.add({
id: 'tt0306646',
title: 'Ringing Bell',
description: 'A baby lamb named Chirin is living an idyllic life on a farm with many other sheep. Chirin is very adventurous and tends to get lost, so he wears a bell around his neck so that his mother can always find him. His mother warns Chirin that he must never venture beyond the fence surrounding the farm, because a huge black wolf lives in the mountains and loves to eat sheep. Chirin is too young and naive to take the advice to heart, until one night the wolf enters the barn and is prepared to kill Chirin, but at the last moment the lamb\'s mother throws herself in the way and is killed instead. The wolf leaves, and Chirin is horrified to see his mother\'s body. Unable to understand why his mother was killed, he becomes very angry and swears that he will go into the mountains and kill the wolf.'
})

ms.add({
id: 'tt1212022',
title: 'The Lion of Judah',
description: 'Follow the adventures of a bold lamb (Judah) and his stable friends as they try to avoid the sacrificial alter the week preceding the crucifixion of Christ. It is a heart-warming account of the Easter story as seen through the eyes of a lovable pig (Horace), a faint-hearted horse (Monty), a pedantic rat (Slink), a rambling rooster (Drake), a motherly cow (Esmay) and a downtrodden donkey (Jack). This magnificent period piece with its epic sets is a roller coaster ride of emotions. Enveloped in humor, this quest follows the animals from the stable in Bethlehem to the great temple in Jerusalem and onto the hillside of Calvary as these unlikely heroes try to save their friend. The journey weaves seamlessly through the biblical accounts of Palm Sunday, Jesus turning the tables in the temple, Peter\'s denial and with a tense, heart-wrenching climax, depicts the crucifixion and resurrection with gentleness and breathtaking beauty. For Judah, the lamb with the heart of a lion, it is a story of courage and faith. For Jack, the disappointed donkey, it becomes a pivotal voyage of hope. For Horace, the, well the dirty pig, and Drake the ignorant rooster, it is an opportunity to do something inappropriate and get into trouble.'
})

it('returns best results for lamb', () => {
// This should be fairly easy. We test that exact matches come before
// prefix matches, and that hits in shorter fields (title) come before
// hits in longer fields (description)
const hits = ms.search('lamb', { fuzzy: 1, prefix: true })
expect(hits.map(({ title }) => title)).toEqual([
// Exact title match.
'Lamb',

// Contains term twice, shortest description.
'Boundin\'',

// Contains term twice.
'Ringing Bell',

// Contains term twice, longest description.
'The Lion of Judah',

// Prefix match in title.
'The Silence of the Lambs'
])
})

it('returns best results for sheep', () => {
// This tests more complex interaction between scoring. We want hits in
// the title to be automatically considered most relevant, because they
// are very short, and the search term occurs less frequently in the
// title than it does in the description. One result, 'Rams', has a very
// short description with an exact match, but it should never outrank
// the result with an exact match in the title AND description.
const hits = ms.search('sheep', { fuzzy: 1, prefix: true })
expect(hits.map(({ title }) => title)).toEqual([
// Has 'sheep' in title and once in a description of average length.
'Shaun the Sheep',

// Has 'sheep' just once, in a short description.
'Rams',

// Contains 'sheep' just once, in a long title.
'Shaun the Sheep: The Farmer\'s Llamas',

// Has most occurrences of 'sheep'.
'Ringing Bell',

// Contains 'sheep' just once, in a long description.
'Lamb'
])
})

it('returns best results for shaun', () => {
// Two movies contain the query in the title. Pick the shorter title.
expect(ms.search('shaun the sheep')[0].title).toEqual('Shaun the Sheep')
expect(ms.search('shaun the sheep', { fuzzy: 1, prefix: true })[0].title).toEqual('Shaun the Sheep')
})

it('returns best results for chirin', () => {
// The title contains neither 'sheep' nor the character name. Movies
// that have 'sheep' or 'the' in the title should not outrank this.
expect(ms.search('chirin the sheep')[0].title).toEqual('Ringing Bell')
expect(ms.search('chirin the sheep', { fuzzy: 1, prefix: true })[0].title).toEqual('Ringing Bell')
})

it('returns best results for judah', () => {
// Title contains the character's name, but the word 'sheep' never
// occurs. Other movies that do contain 'sheep' should not outrank this.
expect(ms.search('judah the sheep')[0].title).toEqual('The Lion of Judah')
expect(ms.search('judah the sheep', { fuzzy: 1, prefix: true })[0].title).toEqual('The Lion of Judah')
})

it('returns best results for bounding', () => {
// The expected hit has an exact match in the description and a fuzzy
// match in the title, and both variations of the term are highly
// specific. Does not contain 'sheep' at all! Because 'sheep' is a
// slightly more common term in the dataset, that should not cause other
// results to outrank this.
expect(ms.search('bounding sheep', { fuzzy: 1 })[0].title).toEqual('Boundin\'')
})
})

describe('song ranking set', () => {
const ms = new MiniSearch({
fields: ['song', 'artist'],
storeFields: ['song']
})

ms.add({
id: '1',
song: 'Killer Queen',
artist: 'Queen'
})

ms.add({
id: '2',
song: 'The Witch Queen Of New Orleans',
artist: 'Redbone'
})

ms.add({
id: '3',
song: 'Waterloo',
artist: 'Abba'
})

ms.add({
id: '4',
song: 'Take A Chance On Me',
artist: 'Abba'
})

ms.add({
id: '5',
song: 'Help',
artist: 'The Beatles'
})

ms.add({
id: '6',
song: 'Yellow Submarine',
artist: 'The Beatles'
})

ms.add({
id: '7',
song: 'Dancing Queen',
artist: 'Abba'
})

ms.add({
id: '8',
song: 'Bohemian Rhapsody',
artist: 'Queen'
})

it('returns best results for witch queen', () => {
const hits = ms.search('witch queen', { fuzzy: 1, prefix: true })
expect(hits.map(({ song }) => song)).toEqual([
// The only result that has both terms. This should not be outranked
// by hits that match only one term.
'The Witch Queen Of New Orleans',

// Contains just one term, but matches both song and artist.
'Killer Queen',

// Match on artist only. Artist is an exact match for 'Queen'.
'Bohemian Rhapsody',

// Match on song only. Song is a worse match for 'Queen'.
'Dancing Queen'
])
})

it('returns best results for queen', () => {
// The only match where both song and artist contain 'queen'.
expect(ms.search('queen', { fuzzy: 1, prefix: true })[0].song).toEqual('Killer Queen')
})
})
})

describe('default tokenization', () => {
Expand Down
Loading