Spaces:
Running on Zero
Running on Zero
File size: 2,538 Bytes
c58c39d 38c5f38 c58c39d 38c5f38 c58c39d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | """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:
log.flush()
print('PARSER_SETUP_ERROR', (RUNTIME/'layout_setup.log').read_text(errors='replace')[-6000:],flush=True)
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:
print('PARSER_EXTRACTION_ERROR', (out/'extraction.log').read_text(errors='replace')[-6000:],flush=True)
raise RuntimeError('LayoutLMv3 could not parse this PDF. Please try a text-based research paper.')
|