BM25 should count doc freq, not field freq

src/main/scala/ai/metarank/feature/matcher/BM25Matcher.scala

@@ -43,12 +43,18 @@ object BM25Matcher {
   case class TermFreqDic(language: String, fields: List[String], docs: Int, avgdl: Double, termfreq: Map[String, Int])
 
   object TermFreqDic {
-    case class Builder(docs: Int = 0, lenSum: Long = 0L, docFreq: Map[String, Int] = Map()) {
-      def withString(terms: Array[String]): Builder = {
+    case class Builder(docs: Int = 0, lenCount: Long = 0L, lenSum: Long = 0L, docFreq: Map[String, Int] = Map()) {
+      def withItem(item: ItemEvent, fields: Set[String], lang: TextAnalyzer): Builder = {
+        val matched = item.fields.flatMap {
+          case StringListField(name, values) if fields.contains(name) => values.map(lang.split)
+          case StringField(name, value) if fields.contains(name)      => List(lang.split(value))
+          case _                                                      => Nil
+        }
         Builder(
           docs = docs + 1,
-          lenSum = lenSum + terms.length,
-          docFreq = terms.foldLeft(docFreq)((acc, next) => {
+          lenSum = lenSum + matched.foldLeft(0L)((acc, terms) => acc + terms.length),
+          lenCount = lenCount + matched.size,
+          docFreq = matched.flatten.distinct.foldLeft(docFreq)((acc, next) => {
             val cnt = acc.getOrElse(next, 0)
             acc.updated(next, cnt + 1)
           })
@@ -59,21 +65,12 @@ object BM25Matcher {
       events
         .through(PrintProgress.tap(None, "events"))
         .collect { case e: ItemEvent => e }
-        .mapChunks(items =>
-          items.flatMap(item =>
-            Chunk.seq(item.fields.flatMap {
-              case StringField(name, value) if fields.contains(name)      => List(value)
-              case StringListField(name, values) if fields.contains(name) => values
-              case _                                                      => Nil
-            })
-          )
-        )
         .compile
-        .fold(Builder())((acc, next) => acc.withString(language.split(next)))
+        .fold(Builder())((acc, next) => acc.withItem(next, fields, language))
         .map(builder =>
           TermFreqDic(
             docs = builder.docs,
-            avgdl = builder.lenSum.toDouble / builder.docs,
+            avgdl = builder.lenSum.toDouble / builder.lenCount,
             termfreq = builder.docFreq,
             language = language.names.head,
             fields = fields.toList.sorted

src/test/scala/ai/metarank/main/TermFreqDicTest.scala

@@ -2,8 +2,10 @@ package ai.metarank.main
 
 import ai.metarank.feature.matcher.BM25Matcher.TermFreqDic
 import ai.metarank.main.command.TermFreq
+import ai.metarank.model.Event
 import ai.metarank.model.Event.ItemEvent
-import ai.metarank.util.RanklensEvents
+import ai.metarank.model.Field.StringField
+import ai.metarank.util.{RanklensEvents, TestItemEvent}
 import better.files.File
 import cats.effect.unsafe.implicits.global
 import org.scalatest.flatspec.AnyFlatSpec
@@ -13,7 +15,7 @@ class TermFreqDicTest extends AnyFlatSpec with Matchers {
   it should "build termfreq from ranklens dataset" in {
     val events = RanklensEvents.stream()
     val tf     = TermFreq.processLanguage("english", Set("title", "description"), events).unsafeRunSync()
-    tf.docs shouldBe 5024
+    tf.docs shouldBe 2512
     tf.termfreq.size shouldBe 11560
   }
 
@@ -30,4 +32,18 @@ class TermFreqDicTest extends AnyFlatSpec with Matchers {
     dir.children.toList.size shouldBe 1
     dir.delete()
   }
+
+  it should "count repetitions only once" in {
+    val event =
+      TestItemEvent("p1", List(StringField("foo", "hello hello world"), StringField("bar", "hello hello world")))
+    val events = fs2.Stream(event)
+    val tf     = TermFreq.processLanguage("english", Set("foo", "bar"), events).unsafeRunSync()
+    tf shouldBe TermFreqDic(
+      language = "en",
+      fields = List("bar", "foo"),
+      docs = 1,
+      avgdl = 3.0,
+      termfreq = Map("hello" -> 1, "world" -> 1)
+    )
+  }
 }