FIRST

Author: chewxy

.travis.yml

@@ -0,0 +1,26 @@
+sudo: false
+language: go
+branches:
+  only:
+    - master
+
+go:
+  - 1.13.x
+  - 1.14.x
+  - tip
+
+env:
+  global:
+    - GOARCH=amd64
+    - TRAVISTEST=true
+
+before_install:
+  - go get github.com/mattn/goveralls
+
+script:
+  - go test -run=. -coverprofile=profile.cov
+  - $HOME/gopath/bin/goveralls -coverprofile=profile.cov -service=travis-ci
+
+matrix:
+  allow_failures:
+    - go: tip

bm25.go

@@ -0,0 +1,73 @@
+// package bm25 is a lingo-friendly BM25 library.
+// BM25 is a scoring function that relies on TFIDF, and is useful for document retrieval
+package bm25
+
+import (
+	"sort"
+
+	"github.com/go-nlp/tfidf"
+	"github.com/xtgo/set"
+)
+
+// DocScore is a tuple of the document ID and a score
+type DocScore struct {
+	ID    int
+	Score float64
+}
+
+// DocScores is a list of DocScore
+type DocScores []DocScore
+
+func (ds DocScores) Len() int           { return len(ds) }
+func (ds DocScores) Less(i, j int) bool { return ds[i].Score < ds[j].Score }
+func (ds DocScores) Swap(i, j int) {
+	ds[i].Score, ds[j].Score = ds[j].Score, ds[i].Score
+	ds[i].ID, ds[j].ID = ds[j].ID, ds[i].ID
+}
+
+// BM25 is the scoring function.
+//
+// k1 should be between 1.2 and 2.
+// b should be around 0.75
+func BM25(tf *tfidf.TFIDF, query tfidf.Document, docs []tfidf.Document, k1, b float64) DocScores {
+	q := tfidf.BOW(query)
+	w := make([]int, len(q))
+	copy(w, q)
+	avgLen := float64(tf.Len) / float64(tf.Docs)
+
+	scores := make([]float64, 0, len(docs))
+	for _, doc := range docs {
+		//TF := tfidf.TF(doc)
+		d := tfidf.BOW(doc)
+		w = append(w, d...)
+		size := set.Inter(sort.IntSlice(w), len(q))
+		n := w[:size]
+
+		score := make([]float64, 0, len(n))
+		docLen := float64(len(d))
+		for _, id := range n {
+			num := (tf.TF[id] * (k1 + 1))
+			denom := (tf.TF[id] + k1*(1-b+b*docLen/avgLen))
+			idf := tf.IDF[id]
+			score = append(score, idf*num/denom)
+		}
+		scores = append(scores, sum(score))
+
+		// reset working vector
+		copy(w, q)
+		w = w[:len(q)]
+	}
+	var retVal DocScores
+	for i := range docs {
+		retVal = append(retVal, DocScore{i, scores[i]})
+	}
+	return retVal
+}
+
+func sum(a []float64) float64 {
+	var retVal float64
+	for _, f := range a {
+		retVal += f
+	}
+	return retVal
+}

bm25_test.go

@@ -0,0 +1,116 @@
+package bm25
+
+import (
+	"fmt"
+	"sort"
+	"strings"
+
+	"github.com/go-nlp/tfidf"
+)
+
+var mobydick = []string{
+	"Call me Ishmael .",
+	"Some years ago -- never mind how long precisely -- having little or no money in my purse , and nothing particular to interest me on shore , I thought I would sail about a little and see the watery part of the world .",
+	"It is a way I have of driving off the spleen and regulating the circulation .",
+	"Whenever I find myself growing grim about the mouth ; ",
+	"whenever it is a damp , drizzly November in my soul ; ",
+	"whenever I find myself involuntarily pausing before coffin warehouses , and bringing up the rear of every funeral I meet ; ",
+	"and especially whenever my hypos get such an upper hand of me , that it requires a strong moral principle to prevent me from deliberately stepping into the street , and methodically knocking people's hats off -- then , I account it high time to get to sea as soon as I can .",
+	"This is my substitute for pistol and ball . ",
+	"With a philosophical flourish Cato throws himself upon his sword ; ",
+	"I quietly take to the ship . There is nothing surprising in this .",
+	"If they but knew it , almost all men in their degree , some time or other , cherish very nearly the same feelings towards the ocean with me .",
+}
+
+type doc []int
+
+func (d doc) IDs() []int { return []int(d) }
+
+func makeCorpus(a []string) (map[string]int, []string) {
+	retVal := make(map[string]int)
+	invRetVal := make([]string, 0)
+	var id int
+	for _, s := range a {
+		for _, f := range strings.Fields(s) {
+			f = strings.ToLower(f)
+			if _, ok := retVal[f]; !ok {
+				retVal[f] = id
+				invRetVal = append(invRetVal, f)
+				id++
+			}
+		}
+	}
+	return retVal, invRetVal
+}
+
+func makeDocuments(a []string, c map[string]int) []tfidf.Document {
+	retVal := make([]tfidf.Document, 0, len(a))
+	for _, s := range a {
+		var ts []int
+		for _, f := range strings.Fields(s) {
+			f = strings.ToLower(f)
+			id := c[f]
+			ts = append(ts, id)
+		}
+		retVal = append(retVal, doc(ts))
+	}
+	return retVal
+}
+
+func Example_BM25() {
+	corpus, _ := makeCorpus(mobydick)
+	docs := makeDocuments(mobydick, corpus)
+	tf := tfidf.New()
+
+	for _, doc := range docs {
+		tf.Add(doc)
+	}
+	tf.CalculateIDF()
+
+	// now we search
+
+	// "ishmael" is a query
+	ishmael := doc{corpus["ishmael"]}
+
+	// "whenever i find" is another query
+	whenever := doc{corpus["whenever"]}
+
+	ishmaelScores := BM25(tf, ishmael, docs, 1.5, 0.75)
+	wheneverScores := BM25(tf, whenever, docs, 1.5, 0.75)
+
+	sort.Sort(sort.Reverse(ishmaelScores))
+	sort.Sort(sort.Reverse(wheneverScores))
+
+	fmt.Printf("Top 3 Relevant Docs to \"Ishmael\":\n")
+	for _, d := range ishmaelScores[:3] {
+		fmt.Printf("\tID   : %d\n\tScore: %1.3f\n\tDoc  : %q\n", d.ID, d.Score, mobydick[d.ID])
+	}
+	fmt.Println("")
+	fmt.Printf("Top 3 Relevant Docs to \"whenever i find\":\n")
+	for _, d := range wheneverScores[:3] {
+		fmt.Printf("\tID   : %d\n\tScore: %1.3f\n\tDoc  : %q\n", d.ID, d.Score, mobydick[d.ID])
+	}
+	// Output:
+	// Top 3 Relevant Docs to "Ishmael":
+	//	ID   : 0
+	//	Score: 3.706
+	//	Doc  : "Call me Ishmael ."
+	//	ID   : 1
+	//	Score: 0.000
+	//	Doc  : "Some years ago -- never mind how long precisely -- having little or no money in my purse , and nothing particular to interest me on shore , I thought I would sail about a little and see the watery part of the world ."
+	//	ID   : 2
+	//	Score: 0.000
+	//	Doc  : "It is a way I have of driving off the spleen and regulating the circulation ."
+	//
+	// Top 3 Relevant Docs to "whenever i find":
+	//	ID   : 3
+	//	Score: 2.031
+	//	Doc  : "Whenever I find myself growing grim about the mouth ; "
+	//	ID   : 4
+	//	Score: 1.982
+	//	Doc  : "whenever it is a damp , drizzly November in my soul ; "
+	//	ID   : 5
+	//	Score: 1.810
+	//	Doc  : "whenever I find myself involuntarily pausing before coffin warehouses , and bringing up the rear of every funeral I meet ; "
+
+}

go.mod

@@ -0,0 +1,8 @@
+module github.com/go-nlp/bm25
+
+go 1.14
+
+require (
+	github.com/go-nlp/tfidf v1.1.0
+	github.com/xtgo/set v1.0.0
+)

go.sum

@@ -0,0 +1,4 @@
+github.com/go-nlp/tfidf v1.0.0 h1:usRGZjJO/MkU4Oq2Xa836MTpAgfhyb2kLyIkjnKfWY0=
+github.com/go-nlp/tfidf v1.0.0/go.mod h1:FHOpf09wrdELx7OnbxywpW4Cs0Q3r15QzpOc73rmUTo=
+github.com/xtgo/set v1.0.0 h1:6BCNBRv3ORNDQ7fyoJXRv+tstJz3m1JVFQErfeZz2pY=
+github.com/xtgo/set v1.0.0/go.mod h1:d3NHzGzSa0NmB2NhFyECA+QdRp29oEn2xbT+TpeFoM8=