mirror of
https://git.pumpin.online/funkpower/catena.git
synced 2026-10-23 03:07:13 +00:00
bindings: the libraries in process, for Python
This commit is contained in:
funkpower
@@ -0,0 +1,235 @@
|
||||
//! The catena libraries behind one surface, for calling in process.
|
||||
//!
|
||||
//! Every tool in this workspace is a library with a thin command on top, and
|
||||
//! a pipeline written in another language can have either. A subprocess is
|
||||
//! the right answer most of the time: it is a process boundary, it composes
|
||||
//! with a shell, and it costs a few milliseconds. In process is the right
|
||||
//! answer when the caller is already holding the text — a stage that would
|
||||
//! otherwise write a file so that it can read it straight back.
|
||||
//!
|
||||
//! What is bound here is the **record shapes**, not the internals. Those
|
||||
//! shapes are documented, versioned, and what every tool already writes to
|
||||
//! stdout; binding the structs behind them would publish field names that are
|
||||
//! free to move and leave the two forms disagreeing about the same chunk.
|
||||
//!
|
||||
//! Nothing here grows logic of its own. If a binding needs something the
|
||||
//! libraries cannot do, that is a missing library function.
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
mod python;
|
||||
|
||||
use corpus_core::{Chunk, Doc, Hit};
|
||||
use cuekit::Cue;
|
||||
|
||||
/// Where chunks are flushed, and how large they may get.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct ChunkBudgets {
|
||||
/// Flush a chunk at this many tokens.
|
||||
pub target: usize,
|
||||
/// Carry this much of each chunk into the next one.
|
||||
pub overlap: usize,
|
||||
/// Drop chunks smaller than this.
|
||||
pub min_tokens: usize,
|
||||
/// Never emit a chunk larger than this.
|
||||
pub hard_max: usize,
|
||||
}
|
||||
|
||||
impl Default for ChunkBudgets {
|
||||
fn default() -> Self {
|
||||
let budgets = chunkmd::Budgets::default();
|
||||
Self {
|
||||
target: budgets.target,
|
||||
overlap: budgets.overlap,
|
||||
min_tokens: budgets.min_tokens,
|
||||
hard_max: budgets.hard_max,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ChunkBudgets> for chunkmd::Budgets {
|
||||
fn from(budgets: ChunkBudgets) -> Self {
|
||||
chunkmd::Budgets {
|
||||
target: budgets.target,
|
||||
overlap: budgets.overlap,
|
||||
min_tokens: budgets.min_tokens,
|
||||
hard_max: budgets.hard_max,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// How much of a recording a passage covers.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct CueBudgets {
|
||||
/// Seconds of audio a passage covers.
|
||||
pub seconds: f64,
|
||||
/// Tokens a passage may hold, whatever the clock says.
|
||||
pub max_tokens: usize,
|
||||
}
|
||||
|
||||
impl Default for CueBudgets {
|
||||
fn default() -> Self {
|
||||
let budgets = cuepack::Budgets::default();
|
||||
Self {
|
||||
seconds: budgets.seconds,
|
||||
max_tokens: budgets.max_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cut markdown into retrievable passages.
|
||||
pub fn chunk_markdown(markdown: &str, budgets: ChunkBudgets) -> Vec<Chunk> {
|
||||
chunkmd::chunk(markdown, &budgets.into())
|
||||
}
|
||||
|
||||
/// Clean converted markdown.
|
||||
pub fn scrub_markdown(markdown: &str) -> String {
|
||||
mdscrub::scrub(markdown)
|
||||
}
|
||||
|
||||
/// Clean converted markdown, normalising page markers as well.
|
||||
///
|
||||
/// Separate because it is not idempotent: running it twice renumbers the
|
||||
/// pages it moved the first time.
|
||||
pub fn scrub_markdown_pages(markdown: &str) -> String {
|
||||
mdscrub::pages::normalize(&mdscrub::scrub(markdown))
|
||||
}
|
||||
|
||||
/// How many tokens a string is, by the encoding every budget here counts in.
|
||||
pub fn count_tokens(text: &str) -> usize {
|
||||
textkit::count(text)
|
||||
}
|
||||
|
||||
/// Read any of the subtitle formats into cues.
|
||||
pub fn parse_subtitles(content: &str) -> Vec<Cue> {
|
||||
cuekit::parse(content)
|
||||
}
|
||||
|
||||
/// Collapse the runs of repeated cues a live caption track leaves behind.
|
||||
pub fn merge_repeats(cues: &[Cue]) -> Vec<Cue> {
|
||||
cuekit::merge_repeats(cues)
|
||||
}
|
||||
|
||||
/// Pack cues into passages worth embedding.
|
||||
pub fn pack_cues(cues: &[Cue], budgets: CueBudgets) -> Vec<Chunk> {
|
||||
cuepack::pack(
|
||||
cues,
|
||||
&cuepack::Budgets {
|
||||
seconds: budgets.seconds,
|
||||
max_tokens: budgets.max_tokens,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Split a collection into the works inside it.
|
||||
///
|
||||
/// `Err` when the document is one work rather than a collection, which is the
|
||||
/// usual answer and not a failure. The message says what each candidate
|
||||
/// produced.
|
||||
pub fn split_collection(markdown: &str) -> Result<Vec<mdsplit::Section>, String> {
|
||||
mdsplit::split(markdown).map_err(|why| why.to_string())
|
||||
}
|
||||
|
||||
/// Fuse ranked lists into one.
|
||||
///
|
||||
/// `weights` is read as channel name to weight; a channel not named is worth
|
||||
/// one. An unknown name is ignored rather than refused, because a caller
|
||||
/// fusing two channels should not have to know about a third.
|
||||
pub fn fuse_hits(
|
||||
hits: &[Hit],
|
||||
k: f64,
|
||||
weights: &[(String, f64)],
|
||||
limit: Option<usize>,
|
||||
) -> Vec<Hit> {
|
||||
let mut channels = std::collections::HashMap::new();
|
||||
for (name, weight) in weights {
|
||||
if let Some(channel) = channel_named(name) {
|
||||
channels.insert(channel, *weight);
|
||||
}
|
||||
}
|
||||
let fused = rrf::fuse(hits, &rrf::Weights { k, channels });
|
||||
match limit {
|
||||
Some(limit) => fused.into_iter().take(limit).collect(),
|
||||
None => fused,
|
||||
}
|
||||
}
|
||||
|
||||
/// The channel a name refers to, spelled as the records spell it.
|
||||
fn channel_named(name: &str) -> Option<corpus_core::Channel> {
|
||||
match name {
|
||||
"vec" => Some(corpus_core::Channel::Vector),
|
||||
"fts" => Some(corpus_core::Channel::Fts),
|
||||
"meta" => Some(corpus_core::Channel::Meta),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Say what a file is from its bytes and its name.
|
||||
pub fn identify(name: &str, bytes: &[u8]) -> Doc {
|
||||
booksniff::sniff(&bytes, name)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const MARKDOWN: &str = "# Chapter One\n\n<!-- Page 41 -->\n\nThe obligation ran both ways, and contemporaries said so plainly whenever they were asked. A lord who could not protect was a lord who could not command.\n";
|
||||
|
||||
#[test]
|
||||
fn the_surface_agrees_with_the_library_underneath() {
|
||||
assert_eq!(
|
||||
chunk_markdown(MARKDOWN, ChunkBudgets::default()),
|
||||
chunkmd::chunk(MARKDOWN, &chunkmd::Budgets::default())
|
||||
);
|
||||
assert_eq!(scrub_markdown(MARKDOWN), mdscrub::scrub(MARKDOWN));
|
||||
assert_eq!(count_tokens("a passage"), textkit::count("a passage"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_default_budgets_are_the_tools_defaults() {
|
||||
let ours = ChunkBudgets::default();
|
||||
let theirs = chunkmd::Budgets::default();
|
||||
assert_eq!(
|
||||
(ours.target, ours.overlap, ours.min_tokens, ours.hard_max),
|
||||
(
|
||||
theirs.target,
|
||||
theirs.overlap,
|
||||
theirs.min_tokens,
|
||||
theirs.hard_max
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
CueBudgets::default().seconds,
|
||||
cuepack::Budgets::default().seconds
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_document_that_is_one_work_says_so_rather_than_splitting() {
|
||||
assert!(split_collection(MARKDOWN).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subtitles_come_back_as_cues_and_pack_into_passages() {
|
||||
let srt = "1\n00:00:01,000 --> 00:00:04,000\nthe mill stood at the northern end\n\n";
|
||||
let cues = parse_subtitles(srt);
|
||||
assert_eq!(cues.len(), 1);
|
||||
let packed = pack_cues(&cues, CueBudgets::default());
|
||||
assert_eq!(packed.len(), 1);
|
||||
assert_eq!(packed[0].start_time, Some(1.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_channel_nobody_here_knows_is_ignored_rather_than_refused() {
|
||||
assert_eq!(channel_named("vec"), Some(corpus_core::Channel::Vector));
|
||||
assert_eq!(channel_named("lexical"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bytes_are_identified_the_same_way_a_file_is() {
|
||||
let pdf = b"%PDF-1.4\n";
|
||||
assert_eq!(
|
||||
identify("book.pdf", pdf).kind,
|
||||
booksniff::sniff(&&pdf[..], "book.pdf").kind
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
//! Python bindings.
|
||||
//!
|
||||
//! These hand back the record shapes as dictionaries — the same JSON every
|
||||
//! tool writes to stdout, spelled the way Python spells things. A stage that
|
||||
//! was reading NDJSON from a subprocess can call the function instead and get
|
||||
//! the dictionaries it was already parsing.
|
||||
//!
|
||||
//! Built only under the `python` feature, which maturin turns on. See
|
||||
//! Cargo.toml for why it cannot be on by default.
|
||||
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyDict, PyList};
|
||||
|
||||
use crate::{ChunkBudgets, CueBudgets};
|
||||
|
||||
/// A record to a Python dictionary, by way of its JSON.
|
||||
///
|
||||
/// Through JSON rather than field by field so that the two forms cannot drift
|
||||
/// apart: if a record gains a field, it appears here without anybody
|
||||
/// remembering to add it.
|
||||
fn to_py<'py, T: serde::Serialize>(py: Python<'py>, value: &T) -> PyResult<Bound<'py, PyAny>> {
|
||||
let json = serde_json::to_value(value)
|
||||
.map_err(|e| PyValueError::new_err(format!("record does not serialise: {e}")))?;
|
||||
json_to_py(py, &json)
|
||||
}
|
||||
|
||||
/// Every record in a list, likewise.
|
||||
fn list_to_py<'py, T: serde::Serialize>(
|
||||
py: Python<'py>,
|
||||
values: &[T],
|
||||
) -> PyResult<Bound<'py, PyList>> {
|
||||
let mut out = Vec::with_capacity(values.len());
|
||||
for value in values {
|
||||
out.push(to_py(py, value)?);
|
||||
}
|
||||
PyList::new(py, out)
|
||||
}
|
||||
|
||||
fn json_to_py<'py>(py: Python<'py>, value: &serde_json::Value) -> PyResult<Bound<'py, PyAny>> {
|
||||
Ok(match value {
|
||||
serde_json::Value::Null => py.None().into_bound(py),
|
||||
serde_json::Value::Bool(b) => b.into_pyobject(py)?.to_owned().into_any(),
|
||||
serde_json::Value::Number(n) => match (n.as_u64(), n.as_i64(), n.as_f64()) {
|
||||
(Some(u), _, _) => u.into_pyobject(py)?.into_any(),
|
||||
(_, Some(i), _) => i.into_pyobject(py)?.into_any(),
|
||||
(_, _, Some(f)) => f.into_pyobject(py)?.into_any(),
|
||||
// JSON cannot hold a number this is not one of, so this is
|
||||
// unreachable rather than a case worth inventing a value for.
|
||||
_ => return Err(PyValueError::new_err(format!("{n} is not a number"))),
|
||||
},
|
||||
serde_json::Value::String(s) => s.into_pyobject(py)?.into_any(),
|
||||
serde_json::Value::Array(items) => {
|
||||
let mut out = Vec::with_capacity(items.len());
|
||||
for item in items {
|
||||
out.push(json_to_py(py, item)?);
|
||||
}
|
||||
PyList::new(py, out)?.into_any()
|
||||
}
|
||||
serde_json::Value::Object(fields) => {
|
||||
let dict = PyDict::new(py);
|
||||
for (key, value) in fields {
|
||||
dict.set_item(key, json_to_py(py, value)?)?;
|
||||
}
|
||||
dict.into_any()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// A Python dictionary back to a record, the same way round.
|
||||
fn from_py<T: serde::de::DeserializeOwned>(value: &Bound<'_, PyAny>) -> PyResult<T> {
|
||||
let json: String = value
|
||||
.py()
|
||||
.import("json")?
|
||||
.call_method1("dumps", (value,))?
|
||||
.extract()?;
|
||||
serde_json::from_str(&json)
|
||||
.map_err(|e| PyValueError::new_err(format!("not a record of the expected shape: {e}")))
|
||||
}
|
||||
|
||||
/// Every dictionary in a sequence.
|
||||
fn from_py_list<T: serde::de::DeserializeOwned>(values: &Bound<'_, PyAny>) -> PyResult<Vec<T>> {
|
||||
let mut out = Vec::new();
|
||||
for value in values.try_iter()? {
|
||||
out.push(from_py(&value?)?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Cut markdown into retrievable passages, one dictionary each.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (markdown, *, target = None, overlap = None, min_tokens = None, hard_max = None))]
|
||||
fn chunk_markdown<'py>(
|
||||
py: Python<'py>,
|
||||
markdown: &str,
|
||||
target: Option<usize>,
|
||||
overlap: Option<usize>,
|
||||
min_tokens: Option<usize>,
|
||||
hard_max: Option<usize>,
|
||||
) -> PyResult<Bound<'py, PyList>> {
|
||||
let default = ChunkBudgets::default();
|
||||
let budgets = ChunkBudgets {
|
||||
target: target.unwrap_or(default.target),
|
||||
overlap: overlap.unwrap_or(default.overlap),
|
||||
min_tokens: min_tokens.unwrap_or(default.min_tokens),
|
||||
hard_max: hard_max.unwrap_or(default.hard_max),
|
||||
};
|
||||
list_to_py(py, &crate::chunk_markdown(markdown, budgets))
|
||||
}
|
||||
|
||||
/// Clean converted markdown.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (markdown, *, pages = false))]
|
||||
fn scrub(markdown: &str, pages: bool) -> String {
|
||||
match pages {
|
||||
true => crate::scrub_markdown_pages(markdown),
|
||||
false => crate::scrub_markdown(markdown),
|
||||
}
|
||||
}
|
||||
|
||||
/// How many tokens a string is.
|
||||
#[pyfunction]
|
||||
fn count_tokens(text: &str) -> usize {
|
||||
crate::count_tokens(text)
|
||||
}
|
||||
|
||||
/// Read any of the subtitle formats into cues.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (content, *, keep_repeats = false))]
|
||||
fn parse_subtitles<'py>(
|
||||
py: Python<'py>,
|
||||
content: &str,
|
||||
keep_repeats: bool,
|
||||
) -> PyResult<Bound<'py, PyList>> {
|
||||
let cues = crate::parse_subtitles(content);
|
||||
let cues = match keep_repeats {
|
||||
true => cues,
|
||||
false => crate::merge_repeats(&cues),
|
||||
};
|
||||
list_to_py(py, &cues)
|
||||
}
|
||||
|
||||
/// Pack cues into passages worth embedding.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (cues, *, seconds = None, max_tokens = None))]
|
||||
fn pack_cues<'py>(
|
||||
py: Python<'py>,
|
||||
cues: &Bound<'py, PyAny>,
|
||||
seconds: Option<f64>,
|
||||
max_tokens: Option<usize>,
|
||||
) -> PyResult<Bound<'py, PyList>> {
|
||||
let default = CueBudgets::default();
|
||||
let budgets = CueBudgets {
|
||||
seconds: seconds.unwrap_or(default.seconds),
|
||||
max_tokens: max_tokens.unwrap_or(default.max_tokens),
|
||||
};
|
||||
let cues: Vec<cuekit::Cue> = from_py_list(cues)?;
|
||||
list_to_py(py, &crate::pack_cues(&cues, budgets))
|
||||
}
|
||||
|
||||
/// Split a collection into the works inside it.
|
||||
///
|
||||
/// Raises `ValueError` when the document is one work rather than a
|
||||
/// collection, with the reason each candidate was rejected.
|
||||
#[pyfunction]
|
||||
fn split_collection<'py>(py: Python<'py>, markdown: &str) -> PyResult<Bound<'py, PyList>> {
|
||||
match crate::split_collection(markdown) {
|
||||
Ok(sections) => list_to_py(py, §ions),
|
||||
Err(why) => Err(PyValueError::new_err(why)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Fuse ranked lists into one.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (hits, *, k = None, weights = None, limit = None))]
|
||||
fn fuse_hits<'py>(
|
||||
py: Python<'py>,
|
||||
hits: &Bound<'py, PyAny>,
|
||||
k: Option<f64>,
|
||||
weights: Option<std::collections::HashMap<String, f64>>,
|
||||
limit: Option<usize>,
|
||||
) -> PyResult<Bound<'py, PyList>> {
|
||||
let hits: Vec<corpus_core::Hit> = from_py_list(hits)?;
|
||||
let weights: Vec<(String, f64)> = weights.unwrap_or_default().into_iter().collect();
|
||||
let fused = crate::fuse_hits(&hits, k.unwrap_or(rrf::K), &weights, limit);
|
||||
list_to_py(py, &fused)
|
||||
}
|
||||
|
||||
/// Say what a file is from its bytes and its name.
|
||||
#[pyfunction]
|
||||
fn identify<'py>(py: Python<'py>, name: &str, data: &[u8]) -> PyResult<Bound<'py, PyAny>> {
|
||||
to_py(py, &crate::identify(name, data))
|
||||
}
|
||||
|
||||
/// The module Python imports.
|
||||
#[pymodule]
|
||||
fn _catena(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
module.add("__version__", env!("CARGO_PKG_VERSION"))?;
|
||||
module.add_function(wrap_pyfunction!(chunk_markdown, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(scrub, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(count_tokens, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(parse_subtitles, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(pack_cues, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(split_collection, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(fuse_hits, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(identify, module)?)?;
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user