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 std::path::{Path, PathBuf};
|
||||||
use text_splitter::{Characters, TextSplitter};
|
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
|
/// Target chunk size handed to the embedding model: small enough that a
|
||||||
/// handful of retrieved chunks stays well within a local model's context
|
/// 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.
|
/// 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>) {
|
fn push_document(path: &Path, splitter: &TextSplitter<Characters>, documents: &mut Vec<Document>) {
|
||||||
match std::fs::read_to_string(path) {
|
match std::fs::read_to_string(path) {
|
||||||
Ok(mut text) => {
|
Ok(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();
|
|
||||||
}
|
|
||||||
let chunks: Vec<&str> = splitter.chunks(&text).collect();
|
let chunks: Vec<&str> = splitter.chunks(&text).collect();
|
||||||
let total = chunks.len();
|
let total = chunks.len();
|
||||||
let source = |i: usize| {
|
let source = |i: usize| {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue