libgrammstein 0.1.0

Hybrid language model (N-gram + Embeddings) for WFST text correction
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

//! RAG (Retrieval-Augmented Generation) module for document retrieval.
//!
//! This module provides dense retrieval capabilities for document search,
//! using ModernBERT embeddings with configurable backends.
//!
//! # Features
//!
//! - **Document Indexing**: Index documents with embeddings and metadata
//! - **Dense Retrieval**: Query documents using embedding similarity
//! - **Exact Backend**: Pre-normalized embeddings with BLAS-accelerated dot product
//! - **HNSW Backend**: Approximate nearest neighbor for large indices (optional)
//! - **Synopsis Generation**: Automatic document summarization
//!
//! # Example
//!
//! ```ignore
//! use libgrammstein::rag::{RagIndex, ExactCosineBackend, Document};
//!
//! let mut index = RagIndex::new(ExactCosineBackend::new(768));
//! index.add_document(document)?;
//!
//! let results = index.query("search query", 10)?;
//! ```

mod backend;
mod builder;
mod document;
mod exact_backend;
#[cfg(feature = "rag-hnsw")]
mod hnsw_backend;
mod index;
mod retriever;

pub use backend::{cosine_similarity, dot_product, BackendType, RetrievalBackend};
pub use builder::{IndexBuilder, IndexBuilderConfig, ParallelIndexBuilder};
pub use document::{
    Document, DocumentBuilder, DocumentId, DocumentMeta, DocumentMetadata, LanguageTag,
};
pub use exact_backend::ExactCosineBackend;
#[cfg(feature = "rag-hnsw")]
pub use hnsw_backend::HnswBackend;
pub use index::{RagIndex, RagIndexConfig};
pub use retriever::{format_results, BatchRetriever, RetrievalConfig, RetrievalResult, Retriever};

// Re-export Synopsis from neural module for convenience
pub use crate::neural::{Synopsis, SynopsisSource};

/// Result type for RAG module operations.
pub type Result<T> = std::result::Result<T, RagError>;

/// Error type for RAG module operations.
#[derive(Debug, thiserror::Error)]
pub enum RagError {
    /// Document not found.
    #[error("Document not found: {0}")]
    DocumentNotFound(DocumentId),

    /// Index error.
    #[error("Index error: {0}")]
    IndexError(String),

    /// Embedding error.
    #[error("Embedding error: {0}")]
    EmbeddingError(String),

    /// I/O error.
    #[error("I/O error: {0}")]
    Io(#[from] std::io::Error),

    /// Serialization error.
    #[error("Serialization error: {0}")]
    Serialization(String),

    /// Neural module error.
    #[error("Neural error: {0}")]
    Neural(#[from] crate::neural::NeuralError),
}