Add document embedding and retrieval via a dedicated embedding model #10
1 changed files with 1 additions and 13 deletions
|
|
@ -2,12 +2,6 @@ use rig::embeddings::{EmbedError, TextEmbedder};
|
|||
use std::path::{Path, PathBuf};
|
||||
use text_splitter::{Characters, TextSplitter};
|
||||
|
||||
/// Safety valve against accidentally pointing `--doc` at a huge binary or
|
||||
/// log file — not a content limit. Real documents are chunked in full (see
|
||||
/// `CHUNK_CHARS` below), so nothing meaningful gets silently dropped short
|
||||
/// of this.
|
||||
const MAX_FILE_CHARS: usize = 2_000_000;
|
||||
|
||||
/// Target chunk size handed to the embedding model: small enough that a
|
||||
/// handful of retrieved chunks stays well within a local model's context
|
||||
/// window, large enough to keep a paragraph or two of context in each one.
|
||||
|
|
@ -61,13 +55,7 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result<Vec<Documen
|
|||
|
||||
fn push_document(path: &Path, splitter: &TextSplitter<Characters>, documents: &mut Vec<Document>) {
|
||||
match std::fs::read_to_string(path) {
|
||||
Ok(mut text) => {
|
||||
// Cheap upper bound on char count, so the vast majority of
|
||||
// documents (well under the limit) skip the char-by-char walk
|
||||
// entirely.
|
||||
if text.len() > MAX_FILE_CHARS && text.chars().count() > MAX_FILE_CHARS {
|
||||
text = text.chars().take(MAX_FILE_CHARS).collect();
|
||||
}
|
||||
Ok(text) => {
|
||||
let chunks: Vec<&str> = splitter.chunks(&text).collect();
|
||||
let total = chunks.len();
|
||||
let source = |i: usize| {
|
||||
|
|
|
|||
Loading…
Reference in a new issue