<commit_before>(ns analyze-data.find-knn
  (:require [clojure.core.matrix :as m]
            [clojure.edn :as edn]
            [clojure.java.io :as io]
            [analyze-data.knn.core :refer [cosine-similarity knn]]
            [analyze-data.tf-idf.core :refer [tf-idf-document to-terms]]
            [analyze-data.gzip :refer [gzip-reader]])
  (:import java.io.PushbackReader))

(defn read-edn-stream
  "Return a lazy sequence of items parsed from a stream of edn objects. reader
  must be a java.io.PushbackReader or some derivee. Since it is lazy, the
  stream must remain open while new items are being requested."
  [reader]
  (lazy-seq
    (when-let [line (edn/read {:eof nil} reader)]
      (cons line (read-edn-stream reader)))))

(defn read-tf-idf-data
  "TODO"
  [file-path]
  (with-open [file-reader (-> file-path
                              gzip-reader
                              PushbackReader.)]
    (let [tf-idf-stream (read-edn-stream file-reader)]
      {:idf (nth tf-idf-stream 0)
       :all-terms (nth tf-idf-stream 1)
       :item-names (nth tf-idf-stream 2)
       :data (m/sparse-matrix (nthrest tf-idf-stream 3))})))

(defn knn-named-result
  "TODO"
  [item-names result]
  (let [index (first result)
        distance (second result)]
    [(nth item-names index) distance]))

(defn find-knn
  "TODO"
  [tf-idf-data query-document]
  (let [{:keys [idf all-terms item-names data]} tf-idf-data
        query-vector (->> query-document
                          (to-terms)
                          (tf-idf-document all-terms idf)
                          (m/sparse-array))
        nearest-neighbors (knn data
                               query-vector
                               :distance-fn cosine-similarity)]
    (map (partial knn-named-result item-names) nearest-neighbors)))<commit_msg>Reduce memory footprint of read-edn-stream.<commit_after>(ns analyze-data.find-knn
  (:require [clojure.core.matrix :as m]
            [clojure.edn :as edn]
            [clojure.java.io :as io]
            [analyze-data.knn.core :refer [cosine-similarity knn]]
            [analyze-data.tf-idf.core :refer [tf-idf-document to-terms]]
            [analyze-data.gzip :refer [gzip-reader]])
  (:import java.io.PushbackReader))

(defn read-edn-stream
  "Return a lazy sequence of (transform-fn item) for each item parsed from a
  stream of edn objects. transform-fn defaults to identity. reader must be a
  java.io.PushbackReader or some derivee. Since it is lazy, the stream must
  remain open while new items are being requested."
  ([reader] (read-edn-stream reader identity))
  ([reader transform-fn]
    (lazy-seq
      (when-let [line (edn/read {:eof nil} reader)]
        (cons (transform-fn line) (read-edn-stream reader transform-fn))))))

(defn read-tf-idf-data
  "TODO"
  [file-path]
  (with-open [file-reader (-> file-path
                              gzip-reader
                              PushbackReader.)]
    {:idf (edn/read file-reader)
     :all-terms (edn/read file-reader)
     :item-names (edn/read file-reader)
     :data (m/sparse-matrix (read-edn-stream file-reader m/sparse-array))}))

(defn knn-named-result
  "TODO"
  [item-names result]
  (let [index (first result)
        distance (second result)]
    [(nth item-names index) distance]))

(defn document-to-vector
  "TODO"
  [all-terms idf document]
  (->> document (to-terms) (tf-idf-document all-terms idf) (m/sparse-array)))

(defn find-knn
  "TODO"
  [tf-idf-data query-document]
  (let [{:keys [idf all-terms item-names data]} tf-idf-data
        query-vector (document-to-vector all-terms idf query-document)
        nearest-neighbors (knn data
                               query-vector
                               :distance-fn cosine-similarity)]
    (map (partial knn-named-result item-names) nearest-neighbors)))
{'fraction_non_alphanumeric': 0.1296972860125261, 'fraction_numerical': 0.0010438413361169101, 'mean_word_length': 2.218303946263644, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3743235', 'n_tokens_mistral': 1323, 'n_tokens_neox': 1296, 'n_words': 323}