TF-IDF in OCaml: When Functional Programming Clicks

I built Cogitator, a local code search engine in OCaml. It indexes code chunks using TF-IDF, ranks results by cosine similarity, and outputs them as RAG prompts for LLMs. The whole thing is about 500 lines of OCaml.

I picked OCaml because I wanted to learn it, and it happened to suit the problem well.

The Tokenization Pipeline

The core of any search engine is tokenization: turning raw text into searchable terms. In Cogitator, the pipeline is:

raw text → split punctuation → split camelCase → split snake_case → lowercase → stem → remove stopwords

In OCaml, this is a pipeline of pure functions composed with |>:

let tokenize text =
  text
  |> split_punctuation
  |> List.concat_map split_camel
  |> List.concat_map split_snake
  |> List.map String.lowercase_ascii
  |> List.filter (fun w -> String.length w > 1)
  |> List.map Stemmer.stem
  |> List.filter (fun w -> not (Stopwords.is_stopword w))

Each function takes a value and returns a value without mutating anything, so you can read the pipeline top to bottom and know what happens at each step. In Python I would've written this as a series of list comprehensions or a loop with intermediate variables, which works but is harder to follow.

The split_camel function uses character-level pattern matching to detect transitions from lowercase to uppercase:

let split_camel s =
  let len = String.length s in
  if len <= 1 then [s]
  else
    let buf = Buffer.create 16 in
    Buffer.add_char buf s.[0];
    let rec go acc i =
      if i >= len then
        let word = Buffer.contents buf in
        if word <> "" then List.rev (word :: acc) else List.rev acc
      else
        let prev = s.[i - 1] in
        let curr = s.[i] in
        if prev >= 'a' && prev <= 'z' && curr >= 'A' && curr <= 'Z' then begin
          let word = Buffer.contents buf in
          Buffer.clear buf;
          Buffer.add_char buf curr;
          go (word :: acc) (i + 1)
        end else begin
          Buffer.add_char buf curr;
          go acc (i + 1)
        end
    in
    go [] 1

"getUserName" becomes ["get"; "User"; "Name"]. The recursive go function with accumulator is idiomatic OCaml: a fold over characters, building the result list as it goes.

TF-IDF: Maps All the Way Down

Most of the TF-IDF computation is folds into OCaml's Map module. Term frequency is a fold over tokens into a StringMap:

let term_frequency tokens =
  let total = float_of_int (List.length tokens) in
  let counts =
    List.fold_left
      (fun acc word ->
        let n = match StringMap.find_opt word acc with
          | Some n -> n + 1
          | None -> 1
        in
        StringMap.add word n acc)
      StringMap.empty tokens
  in
  StringMap.map (fun count -> float_of_int count /. total) counts

Inverse document frequency works the same way. Fold over the corpus, count documents containing each term, take the log:

let inverse_document_frequency corpus =
  let n_docs = float_of_int (List.length corpus) in
  let doc_counts =
    List.fold_left
      (fun acc tokens ->
        let unique = List.sort_uniq String.compare tokens in
        List.fold_left
          (fun acc word ->
            let n = match StringMap.find_opt word acc with
              | Some n -> n + 1
              | None -> 1
            in
            StringMap.add word n acc)
          acc unique)
      StringMap.empty corpus
  in
  StringMap.map (fun df -> log (n_docs /. float_of_int df)) doc_counts

Every intermediate value is immutable. The StringMap.add call returns a new map and leaves the old one untouched. In Python with dictionaries, you'd mutate in place and hope nobody else holds a reference. Here, immutability is the default and the compiler enforces it.

Fuzzy Matching via Edit Distance

Search queries contain typos. Cogitator handles this with Levenshtein edit distance. If a query term isn't in the IDF vocabulary, it finds the closest match within edit distance 2:

let fuzzy_expand query_tokens idf =
  List.fold_left (fun acc qt ->
    if String.length qt < 4 then acc
    else if StringMap.mem qt idf then acc
    else
      let best = List.fold_left (fun (bw, bd) w ->
        let d = edit_distance qt w in
        if d < bd then (Some w, d) else (bw, bd))
        (None, 3) idf_words
      in
      match best with
      | (Some w, d) ->
        let weight = match d with
          | 1 -> 0.5
          | 2 -> 0.25
          | _ -> 0.0
        in
        if weight > 0.0 then StringMap.add w weight acc else acc
      | _ -> acc)
    StringMap.empty query_tokens

The match from distance to weight reads better than an if/else chain would. After a while with pattern matching, languages without it start to feel clumsy.

What OCaml Gave Me

Exhaustive matching

The file_kind type is a variant:

type file_kind =
  | Code of string   (* language name *)
  | Markdown
  | Text

Every match on this type has to handle all three cases, so when I added Markdown as a third variant the compiler flagged every function that didn't handle it. In Python, I'd have found those bugs in production.

Composition

The code has no dependency injection, strategy pattern or abstract base classes. The tokenization pipeline is function composition, and search is tokenize → vectorize → cosine similarity → sort. Each piece can be tested in isolation because there are no side effects to mock.

Refactoring

When I changed the search_result type from a tuple to a record, the compiler caught every usage that needed updating. I didn't need to grep for call sites or run the test suite to find what broke.

The library is about 500 lines and covers tokenization, stemming, TF-IDF indexing, cosine similarity search, fuzzy matching, result merging, cache persistence, and a REPL. The same feature set in Python would be noticeably more code, and I'd be less sure it all fit together.

I wouldn't use OCaml for everything. For a data transformation pipeline with a well-defined domain, where correctness matters, it's a very good fit.