This repository contains a command-line interface (CLI) tool for converting JSON documents outputted by the PDF parsing pipeline to embeddings. The tool is designed to process the parsed documents and generate embeddings using the SentenceTransformer model.
- Python >3.8
- poetry package manager
To run the CLI tool locally, use the following command:
python main.py [OPTIONS] INPUT_DIR OUTPUT_DIR
Alternatively build and use the image via Docker after creating an environment file and filling in the correct variables.
cp .env.example .env
make build
make run_embeddings_generation
--s3
: Specifies whether the input and output directories are in S3. If this option is provided, the tool will read from and write to S3 instead of the local file system.--redo
: Redo encoding for files that have already been parsed. By default, files with IDs that already exist in the output directory are skipped.--device
: Specifies the device to use for embeddings generation. Available options are "cuda" (for GPU) and "cpu".--limit
: Optionally limits the number of text samples to process. Useful for debugging.
INPUT_DIR
: The directory containing the JSON files outputted by the PDF parsing pipeline.OUTPUT_DIR
: The directory to save the embeddings to.
To convert JSON documents located in the input
directory to embeddings and save them in the output
directory, you can use the following command:
python main.py --s3 input output
Make sure to replace input
and output
with the appropriate directories or S3 paths.