File RAG Agent - Offline Doc Search

1#include "bs9/index.hpp"
2 
3#include <algorithm>
4#include <cmath>
5#include <cstring>
6#include <fstream>
7#include <stdexcept>
8#include <unordered_map>
9 
10#include "bs9/text.hpp"
11 
12namespace bs9 {
13namespace {
14 
15constexpr float kBm25K1 = 1.2f;
16constexpr float kBm25B = 0.75f;
17 
18template <typename T>
19void write_pod(std::ostream& os, const T& value) {
20  os.write(reinterpret_cast<const char*>(&value), sizeof(T));
21}
22 
23template <typename T>
24void read_pod(std::istream& is, T& value) {
25  is.read(reinterpret_cast<char*>(&value), sizeof(T));
26  if (!is) throw std::runtime_error("index file truncated");
27}
28 
29void write_string(std::ostream& os, const std::string& s) {
30  const auto len = static_cast<std::uint32_t>(s.size());
31  write_pod(os, len);
32  os.write(s.data(), static_cast<std::streamsize>(len));
33}
34 
35std::string read_string(std::istream& is) {
36  std::uint32_t len = 0;
37  read_pod(is, len);
38  std::string s(len, '\0');
39  is.read(s.data(), static_cast<std::streamsize>(len));
40  if (!is) throw std::runtime_error("index file truncated");
41  return s;
42}
43 
44}  // namespace
45 
46Index::Index(std::shared_ptr<Embedder> embedder)
47    : embedder_(std::move(embedder)), dims_(embedder_->dims()) {}
48 
49void Index::index_postings(std::uint32_t chunk_id,
50                           const std::vector<std::string>& tokens) {
51  std::unordered_map<std::string, std::uint32_t> tf;
52  tf.reserve(tokens.size());
53  for (const auto& t : tokens) {
54    if (is_stop_word(t)) continue;
55    ++tf[t];
56  }
57  for (const auto& [term, count] : tf) {
58    postings_[term].push_back(Posting{chunk_id, count});
59  }
60}
61 
62PreparedDocument prepare_document(const Embedder& embedder, std::string path,
63                                  const std::string& text) {
64  PreparedDocument doc;
65  doc.path = std::move(path);
66 
67  const std::size_t dims = embedder.dims();
68  // Qualified: Index::chunk_text(uint32_t) would otherwise hide the free
69  // function declared in text.hpp.
70  auto chunks = bs9::chunk_text(text);
71 
72  doc.chunks.reserve(chunks.size());
73  doc.texts.reserve(chunks.size());
74  doc.term_freqs.reserve(chunks.size());
75  doc.vectors.resize(chunks.size() * dims);
76 
77  for (std::size_t i = 0; i < chunks.size(); ++i) {
78    const auto tokens = tokenize(chunks[i].text);
79 
80    ChunkMeta meta;
81    meta.offset = chunks[i].offset;
82    meta.length = chunks[i].length;
83    meta.token_count = static_cast<std::uint32_t>(tokens.size());
84    doc.chunks.push_back(meta);
85 
86    embedder.embed(tokens, doc.vectors.data() + i * dims);
87 
88    std::unordered_map<std::string, std::uint32_t> tf;
89    tf.reserve(tokens.size());
90    for (const auto& t : tokens) {
91      if (is_stop_word(t)) continue;
92      ++tf[t];
93    }
94    doc.term_freqs.emplace_back(tf.begin(), tf.end());
95    doc.texts.push_back(std::move(chunks[i].text));
96  }
97 
98  return doc;
99}
100