Spaces:
Running on Zero
Running on Zero
| """Keep the training PDF parser's dependencies separate from the VLM's.""" | |
| import os | |
| import subprocess | |
| import sys | |
| import threading | |
| import venv | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parent | |
| RUNTIME = ROOT/'.runtime' | |
| _lock = threading.Lock() | |
| def worker_environment(): | |
| env = dict(os.environ) | |
| for key in ('HF_TOKEN','HUGGING_FACE_HUB_TOKEN','HF_API_TOKEN','PYTHONPATH','LD_PRELOAD'): | |
| env.pop(key,None) | |
| env.update(HF_HUB_DISABLE_IMPLICIT_TOKEN='1', FORCE_CUDA='0', MAX_JOBS='2', | |
| OMP_NUM_THREADS='2', MKL_NUM_THREADS='2', TOKENIZERS_PARALLELISM='false') | |
| return env | |
| def setup(): | |
| with _lock: | |
| envdir=RUNTIME/'layout_env'; python=envdir/'bin/python' | |
| marker=envdir/'rap_parser_ready' | |
| if marker.exists(): | |
| return str(python) | |
| envdir.parent.mkdir(parents=True,exist_ok=True) | |
| venv.EnvBuilder(with_pip=True,system_site_packages=True).create(envdir) | |
| env=worker_environment() | |
| with (RUNTIME/'layout_setup.log').open('a') as log: | |
| for packages in [ | |
| ['pip','setuptools','wheel','ninja'], | |
| ['transformers==4.44.2','huggingface-hub==0.35.3','numpy<2','opencv-python-headless<4.12'], | |
| ['--no-build-isolation','git+https://github.com/facebookresearch/detectron2.git@1e3e13bbf607b54f62205c4c33922521822fb298'] | |
| ]: | |
| result=subprocess.run([str(python),'-m','pip','install','--disable-pip-version-check',*packages], | |
| env=env,stdout=log,stderr=log,timeout=900) | |
| if result.returncode: | |
| raise RuntimeError('LayoutLMv3 dependency setup failed. Please contact the demo maintainer.') | |
| marker.write_text('layoutlmv3-cpu-v1\n') | |
| return str(python) | |
| def extract(pdf, out): | |
| python=setup() | |
| out=Path(out);out.mkdir(parents=True,exist_ok=True) | |
| with (out/'extraction.log').open('w') as log: | |
| result=subprocess.run([python,str(ROOT/'layout_worker.py'),str(pdf),str(out)], | |
| env=worker_environment(),stdout=log,stderr=log,timeout=900) | |
| if result.returncode: | |
| raise RuntimeError('LayoutLMv3 could not parse this PDF. Please try a text-based research paper.') | |