API Reference

Pipeline orchestration

lit2vec.pipeline.embedding_pipeline(input_dir, output_dir, input_type='parquet', model_name='dunzhang/stella_en_400M_v5', vec_dim=1024, hdf5_name='embedding_v2', group_name='pmids', table_name=None, batch_size=512, overwrite_embeddings=False, overwrite_faiss=False, overwrite_bm25=False, overwrite_years=None, text_field='auto', build_annoy=True, annoy_trees=100, n_results=5, bm25_threshold=12.85, zip_output=False, ui=None)

Run the embedding + indexing pipeline on pre-built parquet or SQL data.

The final output tree matches the structure documented in sample_output/OUTPUT_STRUCTURE.md:

output_dir/
├── config.json
├── <model>.ann                          # monolithic Annoy index
├── bm25/
│   ├── authors/<YYYY>/...
│   └── title_abstract/<YYYY>/...
├── faiss/
│   ├── faiss_quant_<YYYY>_.index
│   └── pmids/embedding_v2_<YYYY>_.parquet
└── metadb_by_year/
    └── pubmed_<YYYY>_.parquet

Intermediate HDF5 embeddings are written to a scratch directory and removed on successful completion so they do not appear in the final tree.

Parameters:

• input_dir (str | Path) – Directory containing year-wise parquet files or a SQL database.

• output_dir (str | Path) – Root directory for the final index bundle.

• input_type (Literal['parquet', 'sql']) – “parquet” or “sql”.

• model_name (str) – HuggingFace sentence-transformer model name.

• vec_dim (int) – Embedding dimensionality (must match the model).

• hdf5_name (str) – Base name for intermediate HDF5 embedding files.

• group_name (str) – HDF5 group name for PMID embeddings.

• table_name (str | None) – Table name when input_type is “sql”.

• batch_size (int) – Embedding batch size.

• overwrite_embeddings (bool) – If True, regenerate all embeddings.

• overwrite_faiss (bool) – If True, rebuild all FAISS indices.

• overwrite_bm25 (bool) – If True, rebuild all BM25 indices.

• overwrite_years (list[int] | None) – Optional list of years to rebuild BM25 for.

• text_field (str) – Text field strategy for parquet input.

• build_annoy (bool) – If True, build the monolithic Annoy index.

• annoy_trees (int) – Number of trees for the Annoy index.

• n_results (int) – Default number of search results (written to config.json).

• bm25_threshold (float) – BM25 score threshold (written to config.json).

• zip_output (bool) – If True, create data.zip inside the output directory.

• ui (PipelineUI | None) – Optional PipelineUI instance for live, step-by-step updates.

lit2vec.pipeline.zip_output_directory(output_dir, archive_name='data.zip')

Create a zip archive of the final pipeline output directory.

Hidden files/directories (.embeddings, .prepared, .pipeline.lock, __pycache__, etc.) and the archive itself are excluded. The archive is written directly inside output_dir.

Parameters:

• output_dir (str | Path) – Root pipeline output directory to archive.

• archive_name (str) – Name of the resulting zip file.

Returns:

Path to the created zip archive.

Return type:

Path

lit2vec.pipeline.export_faiss_pmid_lists(faiss_year_files, hdf5_name, out_dir)

Export per-year Parquet files with a single Int64 ‘pmid’ column.

Output filenames follow the sample-output convention:

<out_dir>/<hdf5_name>_<YYYY>_.parquet

Preserves the HDF5 group enumeration order so that row indices align with the corresponding FAISS index.

Parameters:

• faiss_year_files (dict[int, Path])

• hdf5_name (str)

• out_dir (str | Path)

Return type:

list[str]

lit2vec.pipeline.export_metadb_by_year(year_files, out_dir)

Write metadb_by_year/pubmed_<YYYY>_.parquet with pmid as String.

Preserves source row order so the row index aligns with the BM25 index for that year.

Parameters:

• year_files (dict[int, Path])

• out_dir (str | Path)

Return type:

list[str]

Corpus preparation

Convert a litsync JSONL corpus into year-wise Parquet files for lit2vec.

lit2vec.prepare.litsync.read_manifest(corpus_dir)

Read litsync manifest.json if it exists.

Parameters:

corpus_dir (Path)

Return type:

dict[str, Any]

lit2vec.prepare.litsync.prepare_litsync_corpus(corpus_dir, output_dir, filename_prefix='corpus', progress_callback=None)

Read litsync corpus JSONL files and write year-wise Parquet files.

Returns a summary dict with counts per year and source.

Parameters:

• corpus_dir (str | Path)

• output_dir (str | Path)

• filename_prefix (str)

• progress_callback (Callable[[int, int, str | None], None] | None)

Return type:

dict[str, Any]

lit2vec.prepare.litsync.print_summary_table(summary)

Print a modern Rich summary of the prepared corpus.

Parameters:

summary (dict[str, Any])

Return type:

None

Embeddings

lit2vec.vectordb.create_embeddings.normalize(v, axis=-1, epsilon=1e-08)

Normalizes an ndarray along the specified axis by dividing by its vector norm.

Parameters:

• v (ndarray)

• axis (int)

• epsilon (float)

Return type:

ndarray

lit2vec.vectordb.create_embeddings.embed_query(queries, model, tokenizer, device='cuda')

Embeds a batch of queries using a transformer model and normalizes the embeddings.

Parameters:

• queries (List[str])

• model (transformers.AutoModel)

• tokenizer (transformers.AutoTokenizer)

• device (str)

Return type:

ndarray

lit2vec.vectordb.create_embeddings.save_single_embedding_to_hdf5(hdf5_name, group_name, batch_ids, batch_texts, model, tokenizer, device, batch_size=32)

Saves batch embeddings to an HDF5 file.

Parameters:

• hdf5_name (str)

• group_name (str)

• batch_ids (List[str])

• batch_texts (List[str])

• model (Any)

• tokenizer (Any)

• device (str)

• batch_size (int)

lit2vec.vectordb.create_embeddings.get_model(model_name='dunzhang/stella_en_400M_v5', use_multi_gpu=True, device_map='auto', torch_dtype=torch.float32, low_cpu_mem_usage=True)

Loads the transformer model and tokenizer with optimized GPU usage.

Parameters:

• model_name (str) – Name of the pre-trained model.

• use_multi_gpu (bool) – Whether to use multiple GPUs if available.

• device_map (str) – Device mapping strategy (“auto”, “balanced”, or specific mapping).

• torch_dtype (torch.dtype) – Data type for model weights (float16 saves memory).

• low_cpu_mem_usage (bool) – Whether to use low CPU memory during loading.

Returns:

The loaded transformer model. tokenizer: The corresponding tokenizer. device: The primary device or device map information.

Return type:

model

lit2vec.vectordb.create_embeddings.get_model_with_memory_check(model_name='dunzhang/stella_en_400M_v5', use_multi_gpu=True, max_memory_per_gpu='0.8')

Alternative function that checks GPU memory and loads model accordingly.

Parameters:

• model_name (str) – Name of the pre-trained model.

• use_multi_gpu (bool) – Whether to use multiple GPUs if available.

• max_memory_per_gpu (str) – Maximum memory to use per GPU (e.g., “0.8” for 80%).

Returns:

model, tokenizer, device_info

lit2vec.vectordb.create_embeddings.load_data_from_sql(filename, table_name, columns, limited_rows=100)

Load data from a DuckDB/SQLite database and return a pandas DataFrame.

Parameters:

• filename (str) – Path to the database file.

• table_name (str) – Name of the table to query.

• columns (List[str]) – List of column names to retrieve.

• limited_rows (Optional[int]) – Maximum number of rows to fetch (default: 100). Pass None for no limit.

Returns:

The resulting DataFrame with the queried data.

Return type:

pd.DataFrame

lit2vec.vectordb.create_embeddings.from_sql(db_file, hdf5_name, hdf5_path=None, group_name='pmids', table_name='PubmedDB', model_name='dunzhang/stella_en_400M_v5', batch_size=32, columns=['pmid', 'abstract', 'title', 'authors'], progress_callback=None)

Parameters:

• db_file (str)

• hdf5_name (str)

• hdf5_path (str | None)

• group_name (str)

• table_name (str)

• model_name (str)

• batch_size (int)

• columns (List[str])

• progress_callback (Callable[[int, int, str | None], None] | None)

lit2vec.vectordb.create_embeddings.from_parquet(db_path, hdf5_name, hdf5_path, group_name='pmids', model_name='dunzhang/stella_en_400M_v5', batch_size=1024, overwrite=True, text_field='auto', years=None, progress_callback=None)

Stream Parquet rows in batches to minimize memory usage.

Reads only the required columns per batch using PyArrow and processes each batch without materializing the whole file.

Parameters:

• text_field (str) – Strategy for selecting text to embed. One of: - “auto”: abstract if non-empty, else body (default) - “title”: title only - “abstract”: abstract only - “body”: body/full-text only - “title+abstract”: title and abstract concatenated - “title+abstract+body”: title, abstract, and body concatenated

• years (List[str] | None) – Optional list of year labels (e.g. [“2024”]) to restrict processing to. When given, parquet files for other years are skipped entirely (important for efficient resume: completed years are not re-embedded on the GPU only to be skipped at write time).

• db_path (str)

• hdf5_name (str)

• hdf5_path (str)

• group_name (str)

• model_name (str)

• batch_size (int)

• overwrite (bool)

• progress_callback (Callable[[int, int, str | None], None] | None)

lit2vec.vectordb.create_embeddings.generate_embeddings_core(input_type, **kwargs)

Parameters:

input_type (str)

FAISS index

lit2vec.vectordb.create_faiss_index.create_compressed_index(hdf5_paths, faiss_path)

Creates compund index HNSW and IVF

Parameters:

• hdf5_paths (List[str | Path])

• faiss_path (str | Path)

lit2vec.vectordb.create_faiss_index.create_full_index(hdf5_paths, faiss_path)

Creates uncompressed HNSW index: SOTA

Parameters:

• hdf5_paths (List[str | Path])

• faiss_path (str | Path)

lit2vec.vectordb.create_faiss_index.create_full_quantized_index(hdf5_paths, faiss_path, progress_callback=None)

float32 is quantized to 4bits - saves memory at the cost of speed

Parameters:

• hdf5_paths (List[str | Path])

• faiss_path (str | Path)

• progress_callback (Callable[[int, int, str | None], None] | None)

Annoy index

Build a monolithic Annoy index from per-year HDF5 embedding files.

lit2vec.vectordb.create_annoy_index.build_annoy_index(hdf5_paths, output_path, dim=1024, n_trees=100, metric='angular', progress_callback=None)

Build a single Annoy index from all per-year HDF5 embeddings.

Parameters:

• hdf5_paths (List[str | Path]) – List of paths to yearly HDF5 embedding files. Each file is expected to contain a pmids group with one subgroup per PMID, and each subgroup must have an abstract dataset holding the embedding vector.

• output_path (str | Path) – Destination path for the Annoy index (e.g. stella_en_400M_v5.ann).

• dim (int) – Vector dimensionality. Must match the embedding model output.

• n_trees (int) – Number of Annoy trees. More trees improve recall at the cost of build time and index size.

• metric (str) – Annoy distance metric. angular is appropriate for cosine similarity on L2-normalized vectors.

• progress_callback (Callable[[int, int, str | None], None] | None) – Optional callback(current, total, message).

Returns:

Dict with index_path, total_vectors, dim, n_trees, metric.

Return type:

dict

BM25 index

lit2vec.bm25_index.create_index.index_title_and_abstract_core(file_list, out_dir, progress_callback=None)

Parameters:

• file_list (List[str | Path])

• out_dir (Path | str)

• progress_callback (Callable[[int, int, str | None], None] | None)

lit2vec.bm25_index.create_index.index_authors_core(file_list, out_dir, progress_callback=None)

Parameters:

• file_list (List[str | Path])

• out_dir (Path | str)

• progress_callback (Callable[[int, int, str | None], None] | None)

Metadata databases

lit2vec.metadatadb.create_metadata_sqllite.build_(file, out_dir, columns_sql, columns, values, primary_key)

Parameters:

• file (Path)

• out_dir (Path)

• columns_sql (str)

• columns (str)

• values (str)

• primary_key (str)

lit2vec.metadatadb.create_metadata_sqllite.build_sqllitedb_core(file_list, out_dir, primary_key, progress_callback=None)

Parameters:

• file_list (List[str | Path])

• out_dir (Path | str)

• primary_key (str)

• progress_callback (Callable[[int, int, str | None], None] | None)

Return type:

None

Verification (proctor)

Proctor agent — verify Lit2Vec pipeline outputs for integrity and consistency.

This module validates the final output tree produced by lit2vec run (with or without --zip). It checks that every artifact is readable, non-corrupted, internally consistent, and aligned across the different index types.

exception lit2vec.proctor.ProctorError

Bases: Exception

Raised when a proctor check fails.

class lit2vec.proctor.CheckResult

Bases: object

CheckResult(name: ‘str’, passed: ‘bool’, message: ‘str’ = ‘’, details: ‘dict[str, Any]’ = <factory>, severity: ‘str’ = ‘error’)

name: str

passed: bool

message: str = ''

details: dict[str, Any]

severity: str = 'error'

__init__(name, passed, message='', details=<factory>, severity='error')

Parameters:

• name (str)

• passed (bool)

• message (str)

• details (dict[str, Any])

• severity (str)

Return type:

None

class lit2vec.proctor.ProctorReport

Bases: object

ProctorReport(output_dir: ‘Path’, checks: ‘list[CheckResult]’ = <factory>)

output_dir: Path

checks: list[CheckResult]

property passed: bool

property failed: list[CheckResult]

add(name, passed, message='', severity='error', **details)

Parameters:

• name (str)

• passed (bool)

• message (str)

• severity (str)

• details (Any)

Return type:

None

raise_for_failure()

Return type:

None

__init__(output_dir, checks=<factory>)

Parameters:

• output_dir (Path)

• checks (list[CheckResult])

Return type:

None

lit2vec.proctor.proctor_pipeline_output(output_dir, *, expect_zip=False, model_name='dunzhang/stella_en_400M_v5', vec_dim=1024)

Run the full proctor suite against a Lit2Vec pipeline output directory.

Parameters:

• output_dir (str | Path) – Root directory containing the pipeline artifacts.

• expect_zip (bool) – If True, require a valid data.zip archive in the root.

• model_name (str) – Model name used to derive the expected Annoy filename.

• vec_dim (int) – Expected vector dimensionality.

Returns:

A ProctorReport with one CheckResult per validation step.

Return type:

ProctorReport

lit2vec.proctor.print_report(report)

Print a human-readable summary of a proctor report.

Parameters:

report (ProctorReport)

Return type:

None

Utilities

lit2vec.utils.find_files(dir, prefix=None, suffix=None, file_format=None, recursive=False, include_patterns=None, exclude_patterns=None, min_size_bytes=None, max_files=None, sort_by='name', reverse=False)

Find files in a directory with flexible filtering and sorting.

Parameters:

• dir (str | Path) – Directory to search.

• prefix (str | None) – Filename should start with this value (case-sensitive). If None, ignored.

• suffix (str | None) – Filename should end with this value (before extension). If None, ignored.

• file_format (str | None) – Required extension without dot (e.g., “parquet”). If None, any extension allowed.

• recursive (bool) – If True, search subdirectories.

• include_patterns (list[str] | None) – Additional fnmatch-style patterns (match if any). Applied on basename.

• exclude_patterns (list[str] | None) – Fnmatch patterns to exclude (exclude if any matches). Applied on basename.

• min_size_bytes (int | None) – Keep files whose size >= this value.

• max_files (int | None) – If set, return at most this many files after sorting.

• sort_by (str) – Sort key: “name” (default), “mtime” (modification time), or “size” (bytes).

• reverse (bool) – Reverse the sort order.

Returns:

FindFilesResult with matching file paths and metadata.

Raises:

• FileNotFoundError – If the directory does not exist.

• NotADirectoryError – If the path is not a directory.

• ValueError – If sort_by is invalid.

Return type:

FindFilesResult

lit2vec.utils.sort_files_by_year(files, sort_by='name')

Group files by 4-digit year inferred from filename and return a dict mapping year->list[Path].

Extraction priority: - ‘pubmed_sorted_YYYY’ - ‘pubmed_YYYY’ - any 4-digit year (1900-2099) in the basename

Parameters:

• files (list[Path]) – List of Path-like items.

• sort_by (str) – How to sort files within each year: ‘name’ | ‘mtime’ | ‘size’.

Returns:

mapping from year to list of files.

Return type:

dict[int, list[Path]]

lit2vec.utils.attach_function_as_method(target, func, name=None, method_kind='instance')

Attach a standalone function to an object or class as a method at runtime.

Parameters:

• target (Any) – An instance or a class to which the function will be attached.

• func (Callable) – The function to attach.

• name (str | None) – The attribute name under which to attach the function. Defaults to func.__name__.

• method_kind (str) – One of {“instance”, “class”, “static”}.

Return type:

str

Behavior:

• instance: binds to a single instance (if target is an instance) using MethodType;

  if target is a class, sets an unbound function on the class so future instances receive it as a normal method.

• class: attaches as a classmethod on the class.

• static: attaches as a staticmethod on the class.

Returns:

The final attribute name used.

Parameters:

• target (Any)

• func (Callable)

• name (str | None)

• method_kind (str)

Return type:

str

lit2vec.utils.attach_methods_bulk(target, funcs, method_kind='instance')

Attach multiple functions to a target as methods.

Parameters:

• target (Any) – instance or class.

• funcs (dict[str, Callable] | list[Callable]) – mapping of name->callable or list of callables (names from __name__).

• method_kind (str) – ‘instance’ | ‘class’ | ‘static’.

Returns:

List of attribute names attached.

Return type:

list[str]

lit2vec.utils.search_files(root, pattern=None, regex=None, case_insensitive=True, max_results=None)

Recursively search for files under ‘root’. - pattern: glob pattern (fast). If provided, uses rglob. - regex: match against full path string. - If both provided, results must satisfy BOTH.

Parameters:

• root (str | Path)

• pattern (str | None)

• regex (str | None)

• case_insensitive (bool)

• max_results (int | None)

Return type:

List[Path]

lit2vec.utils.find_one(root, pattern=None, regex=None, case_insensitive=True)

Parameters:

• root (str | Path)

• pattern (str | None)

• regex (str | None)

• case_insensitive (bool)

Return type:

Path | None