clarinpl/embed_documents repository overview

⁠Script for embedding with transformer models

Use this script to embed documents using models based on huggingface transformers.

⁠Running script via docker

Let's assume that $(pwd)/input contains files that you would like to embedd and $(pwd)/output is a folder where result will be stored

!! IMPORTANT !! Please make sure that your output directory exists before running the command below. Otherwise this directory will be created by docker deamon, which will give it root ownership by default!

docker run \
    -it \
    -u "$(id -u):$(id -g)" \
    -v $(pwd)/input:/workspace/input \
    -v $(pwd)/output:/workspace/output \
    clarinpl/embed_documents \
    --input_dir /workspace/input \
    --output /workspace/output/test.json \
    --batch_size 8 \
    --model clarin-pl/herbert-kgr10 \
    --pooling cls

⁠Explanation

• --input_dir: Directory containing documents for which you woul like to get embedding vectors
• --output: Path to output json file
• --batch_size: Processing batch size. Increasing batch_size will incerease your gpu utilisation, but it will consume more VRAM. If you get out of memory error, consider lowering your batch size.
• --model: A model name. Can be either a local path or name from https://huggingface.co/models⁠
• --pooling: What type of pooling you would like to use, to get a single vector representing whole document. Can be one of: max, cls, mean. Usually cls or mean is prefered.

⁠Running via docker with GPU

docker run \
    --gpus device=0 \
    -it \
    -u "$(id -u):$(id -g)" \
    -v $(pwd)/input:/workspace/input \
    -v $(pwd)/output:/workspace/output \
    clarinpl/embed_documents \
    --input_dir /workspace/input \
    --output /workspace/output/test.json \
    --batch_size 8 \
    --model clarin-pl/herbert-kgr10 \
    --pooling cls

where --gpus device=0 defines id of gpu that will be used. When --gpus is not specified, CPU will be used for embedding

⁠Running via locker with local model cache

If you don't want to download model every time you run a script via docker, you can mount your local huggingface cache with -v ${HOME}/cache/huggingface:/.cache/huggingface, eg.:

docker run \
    -it \
    -u "$(id -u):$(id -g)" \
    -v $(pwd)/input:/workspace/input \
    -v $(pwd)/output:/workspace/output \
    -v ()
    clarinpl/embed_documents \
    --input_dir /workspace/input \
    --output /workspace/output/test.json \
    --batch_size 8 \
    --model clarin-pl/herbert-kgr10 \
    --pooling cls

⁠Running without docker

Example is analogous to the one presented in section Running via docker with GPU

CUDA_VISIBLE_DEVICES=0 python3 run_script.py \
    --input_dir input \
    --output test.json \
    --batch_size 32 \
    --model clarin-pl/herbert-kgr10 \
    --pooling cls

If you don't want to use any GPU, just replace CUDA_VISIBLE_DEVICES=0 with CUDA_VISIBLE_DEVICES=-1

⁠Final notes

⁠Authors:

• Michał Swędrowski ([email protected]⁠)
• Michał Pogoda ([email protected]⁠)

Please feel free to contact us when you need any help or information :)

⁠Known Limitations

For now, the script can only use single GPU for embedding. If you need support for multi-gpu, please write to [email protected]⁠