#!/bin/bash

DOCKER_NAME=pytorch_rocm
docker stop ${DOCKER_NAME} && docker rm ${DOCKER_NAME}

# --rm -it
docker run -d -ti --name=${DOCKER_NAME} \
  --device=/dev/kfd \
  --device=/dev/dri \
  --security-opt seccomp=unconfined \
  --shm-size=8G \
  -v $HOME:/workspace \
  rocm/pytorch:rocm7.2_ubuntu24.04_py3.12_pytorch_release_2.9.1

sleep 5
# konteineri sees 
docker exec -ti -u 0 ${DOCKER_NAME} bash -c "
echo 'TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1' >> /root/.bashrc;
source /root/.bashrc;
pip install --upgrade pip;
pip install --upgrade pip;
pip install marker-pdf;
"
# Kontrolli GPU kasutamist konteineri sees:
docker exec -ti -u 0 ${DOCKER_NAME} python3 -c "import torch; print(torch.cuda.is_available())"
# Peaks väljastama `True`
docker exec -ti -u 0 ${DOCKER_NAME} python3 -c "import torch; print(torch.cuda.get_device_name(0))"

# GPU kontrollimine (kui käivitada midagi docker'i sees, siis on näha)
amd-smi

# Dockeri parameetrite selgitus
#   - `--rm`: Kustutab konteineri pärast väljumist.
#   - `-it`: Interaktiivne terminal.
#   - `--device=/dev/kfd --device=/dev/dri`: Annab juurdepääsu AMD GPU seadmetele ROCm-i jaoks.
#   - `--security-opt seccomp=unconfined`: Vajalik mõnede ROCm funktsioonide jaoks (nt HSA).
#   - `--shm-size=8G`: Suurendab jagatud mälu (kasulik suurte mudelite jaoks, nt AI-tööd).
#   - `-v $HOME:/workspace`: Ühendab sinu kodukataloogi konteineri sisse, et saaksid faile jagada.

docker exec -ti -u 0 ${DOCKER_NAME} bash -c "
input_dir=/workspace/rag-demo/pdf-pipeline/data/raw_pdfs \
output_dir=/workspace/rag-demo/pdf-pipeline/data/markdown_output \
cd \"$input_dir\" \
for file in *.pdf; do \
    base=\"${file%.pdf}\" \
    outdir=\"$output_dir\"/\"$base\" \
    if [ -d "$outdir" ]; then \
        echo \"Jätan vahele (kaust olemas): $file\" \
        continue \
    fi \
        echo \"Töötlen pdf faili: $file\" \
    marker_single \"$file\" --output_dir \"$output_dir\" \
done
"