- )({PIPE_REGEX_BASE})(?=)
+ """
- def get_relative_link(url):
+ def replace_component(match):
+ name = match.group()
+ preceding = output[match.start(0) - 50 : match.start(0)]
+ if (
+ "DEFAULT:"
+ not in preceding
+ # and output[: match.start(0)].count("")
+ # > output[match.end(0) :].count("")
+ ):
+ try:
+ ep_url = factories_entry_points[name]
+ except KeyError:
+ pass
+ else:
+ if ep_url.split("#")[0].strip("/") != page.file.url.strip("/"):
+ return "{name}".format(href=ep_url, name=name)
+ return name
+
+ def replace_link(match):
+ relative_url = url = match.group(1) or match.group(2) or match.group(3)
page_url = os.path.join("/", page.file.url)
if url.startswith("/"):
- url = os.path.relpath(url, page_url)
- return url
-
- def replace_component_span(span):
- content = span.text
- if content is None:
- return
- link_url = get_component_url(content.strip("\"'"))
- if link_url is None:
- return
- a = etree.Element("a", href="/" + link_url)
- a.text = content
- span.text = ""
- span.append(a)
-
- def replace_component_names(root):
- # Iterate through all span elements
- spans = list(root.iter("span", "code"))
- for i, span in enumerate(spans):
- prev = span.getprevious()
- if span.getparent().tag == "a":
- continue
- # To avoid replacing default component name in parameter tables
- if prev is None or prev.text != "DEFAULT:":
- replace_component_span(span)
- # if span.text == "add_pipe":
- # next_span = span.getnext()
- # if next_span is None:
- # continue
- # next_span = next_span.getnext()
- # if next_span is None or next_span.tag != "span":
- # continue
- # replace_component_span(next_span)
- # continue
- # tokens = ["@", "factory", "="]
- # while True:
- # if len(tokens) == 0:
- # break
- # if span.text != tokens[0]:
- # break
- # tokens = tokens[1:]
- # span = span.getnext()
- # while span is not None and (
- # span.text is None or not span.text.strip()
- # ):
- # span = span.getnext()
- # if len(tokens) == 0:
- # replace_component_span(span)
-
- # Convert the modified tree back to a string
- return root
-
- def replace_absolute_links(root):
- # Iterate through all a elements
- for a in root.iter("a"):
- href = a.get("href")
- if href is None or href.startswith("http"):
- continue
- a.set("href", get_relative_link(href))
- for img in root.iter("img"):
- href = img.get("src")
- if href is None or href.startswith("http"):
- continue
- img.set("src", get_relative_link(href))
-
- # Convert the modified tree back to a string
- return root
+ relative_url = os.path.relpath(url, page_url)
+ return f'"{relative_url}"'
# Replace absolute paths with path relative to the rendered page
- from lxml.html import etree
-
- root = etree.HTML(output)
- root = replace_component_names(root)
- root = replace_absolute_links(root)
- doctype = root.getroottree().docinfo.doctype
- res = etree.tostring(root, encoding="unicode", method="html", doctype=doctype)
- return res
+ output = regex.sub(PIPE_REGEX, replace_component, output)
+ output = regex.sub(HREF_REGEX, replace_link, output)
+
+ return output
diff --git a/docs/trainable-pipes.md b/docs/trainable-pipes.md
index b9dc02c8..ae6f37b0 100644
--- a/docs/trainable-pipes.md
+++ b/docs/trainable-pipes.md
@@ -97,12 +97,12 @@ class MyComponent(TrainablePipe):
"my-feature": ...(doc),
}
- def collate(self, batch, device: torch.device) -> Dict:
+ def collate(self, batch) -> Dict:
# Collate the features of the "embedding" subcomponent
# and the features of this component as well
return {
- "embedding": self.embedding.collate(batch["embedding"], device),
- "my-feature": torch.as_tensor(batch["my-feature"], device=device),
+ "embedding": self.embedding.collate(batch["embedding"]),
+ "my-feature": torch.as_tensor(batch["my-feature"]),
}
def forward(self, batch: Dict, supervision=False) -> Dict:
diff --git a/edspdf/__init__.py b/edspdf/__init__.py
index 9d5ed750..1ab0d6f1 100644
--- a/edspdf/__init__.py
+++ b/edspdf/__init__.py
@@ -3,7 +3,8 @@
from .pipeline import Pipeline, load
from .registry import registry
from .structures import Box, Page, PDFDoc, Text, TextBox, TextProperties
+from . import data
from . import utils # isort:skip
-__version__ = "0.8.1"
+__version__ = "0.9.0"
diff --git a/edspdf/accelerators/base.py b/edspdf/accelerators/base.py
index 07bc01b2..2360377d 100644
--- a/edspdf/accelerators/base.py
+++ b/edspdf/accelerators/base.py
@@ -1,97 +1,2 @@
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Union
-
-from ..structures import PDFDoc
-
-
-class FromDictFieldsToDoc:
- def __init__(self, content_field, id_field=None):
- self.content_field = content_field
- self.id_field = id_field
-
- def __call__(self, item):
- if isinstance(item, dict):
- return PDFDoc(
- content=item[self.content_field],
- id=item[self.id_field] if self.id_field else None,
- )
- return item
-
-
-class ToDoc:
- @classmethod
- def __get_validators__(cls):
- yield cls.validate
-
- @classmethod
- def validate(cls, value, config=None):
- if isinstance(value, str):
- value = {"content_field": value}
- if isinstance(value, dict):
- value = FromDictFieldsToDoc(**value)
- if callable(value):
- return value
- raise TypeError(
- f"Invalid entry {value} ({type(value)}) for ToDoc, "
- f"expected string, a dict or a callable."
- )
-
-
-FROM_DOC_TO_DICT_FIELDS_TEMPLATE = """
-def fn(doc):
- return {X}
-"""
-
-
-class FromDocToDictFields:
- def __init__(self, mapping):
- self.mapping = mapping
- dict_fields = ", ".join(f"{repr(k)}: doc.{v}" for k, v in mapping.items())
- local_vars = {}
- exec(FROM_DOC_TO_DICT_FIELDS_TEMPLATE.replace("X", dict_fields), local_vars)
- self.fn = local_vars["fn"]
-
- def __reduce__(self):
- return FromDocToDictFields, (self.mapping,)
-
- def __call__(self, doc):
- return self.fn(doc)
-
-
-class FromDoc:
- """
- A FromDoc converter (from a PDFDoc to an arbitrary type) can be either:
-
- - a dict mapping field names to doc attributes
- - a callable that takes a PDFDoc and returns an arbitrary type
- """
-
- @classmethod
- def __get_validators__(cls):
- yield cls.validate
-
- @classmethod
- def validate(cls, value, config=None):
- if isinstance(value, dict):
- value = FromDocToDictFields(value)
- if callable(value):
- return value
- raise TypeError(
- f"Invalid entry {value} ({type(value)}) for ToDoc, "
- f"expected dict or callable"
- )
-
-
class Accelerator:
- def __call__(
- self,
- inputs: Iterable[Any],
- model: Any,
- to_doc: ToDoc = FromDictFieldsToDoc("content"),
- from_doc: FromDoc = lambda doc: doc,
- ):
- raise NotImplementedError()
-
-
-if TYPE_CHECKING:
- ToDoc = Union[str, Dict[str, Any], Callable[[Any], PDFDoc]] # noqa: F811
- FromDoc = Union[Dict[str, Any], Callable[[PDFDoc], Any]] # noqa: F811
+ pass
diff --git a/edspdf/accelerators/multiprocessing.py b/edspdf/accelerators/multiprocessing.py
index 41ef2f80..b11d556e 100644
--- a/edspdf/accelerators/multiprocessing.py
+++ b/edspdf/accelerators/multiprocessing.py
@@ -1,338 +1,15 @@
-import gc
-import signal
-from multiprocessing.connection import wait
-from random import shuffle
-from typing import Any, Iterable, List, Optional, Union
+from typing import List, Optional, Union
import torch
-import torch.multiprocessing as mp
-from .. import TrainablePipe
from ..registry import registry
-from ..utils.collections import batchify
-from .base import Accelerator, FromDictFieldsToDoc, FromDoc, ToDoc
-
-DEBUG = False
-
-debug = (
- (lambda *args, flush=False, **kwargs: print(*args, **kwargs, flush=True))
- if DEBUG
- else lambda *args, **kwargs: None
-)
-
-
-class Exchanger:
- def __init__(
- self,
- num_stages,
- num_gpu_workers,
- num_cpu_workers,
- gpu_worker_devices,
- ):
- # queue for cpu input tasks
- self.gpu_worker_devices = gpu_worker_devices
- # We add prioritized queue at the end for STOP signals
- self.cpu_inputs_queues = [
- [mp.SimpleQueue()] + [mp.SimpleQueue() for _ in range(num_stages + 1)]
- # The input queue is not shared between processes, since calling `wait`
- # on a queue reader from multiple processes may lead to a deadlock
- for _ in range(num_cpu_workers)
- ]
- self.gpu_inputs_queues = [
- [mp.SimpleQueue() for _ in range(num_stages + 1)]
- for _ in range(num_gpu_workers)
- ]
- self.outputs_queue = mp.Queue()
-
- def get_cpu_tasks(self, idx):
- while True:
- queue_readers = wait(
- [queue._reader for queue in self.cpu_inputs_queues[idx]]
- )
- stage, queue = next(
- (stage, q)
- for stage, q in reversed(list(enumerate(self.cpu_inputs_queues[idx])))
- if q._reader in queue_readers
- )
- try:
- item = queue.get()
- except BaseException:
- continue
- if item is None:
- return
- yield stage, item
-
- def put_cpu(self, item, stage, idx):
- return self.cpu_inputs_queues[idx][stage].put(item)
-
- def get_gpu_tasks(self, idx):
- while True:
- queue_readers = wait(
- [queue._reader for queue in self.gpu_inputs_queues[idx]]
- )
- stage, queue = next(
- (stage, q)
- for stage, q in reversed(list(enumerate(self.gpu_inputs_queues[idx])))
- if q._reader in queue_readers
- )
- try:
- item = queue.get()
- except BaseException: # pragma: no cover
- continue
- if item is None:
- return
- yield stage, item
-
- def put_gpu(self, item, stage, idx):
- return self.gpu_inputs_queues[idx][stage].put(item)
-
- def put_results(self, items):
- self.outputs_queue.put(items)
-
- def iter_results(self):
- for out in iter(self.outputs_queue.get, None):
- yield out
-
-
-class CPUWorker(mp.Process):
- def __init__(
- self,
- cpu_idx: int,
- exchanger: Exchanger,
- gpu_pipe_names: List[str],
- model: Any,
- device: Union[str, torch.device],
- ):
- super(CPUWorker, self).__init__()
-
- self.cpu_idx = cpu_idx
- self.exchanger = exchanger
- self.gpu_pipe_names = gpu_pipe_names
- self.model = model
- self.device = device
-
- def _run(self):
- # Cannot pass torch tensor during init i think ? otherwise i get
- # ValueError: bad value(s) in fds_to_keep
- mp._prctl_pr_set_pdeathsig(signal.SIGINT)
-
- model = self.model.to(self.device)
- stages = [{"cpu_components": [], "gpu_component": None}]
- for name, component in model.pipeline:
- if name in self.gpu_pipe_names:
- stages[-1]["gpu_component"] = component
- stages.append({"cpu_components": [], "gpu_component": None})
- else:
- stages[-1]["cpu_components"].append(component)
-
- next_batch_id = 0
- active_batches = {}
- debug(
- f"CPU worker {self.cpu_idx} is ready",
- next(model.parameters()).device,
- flush=True,
- )
-
- had_error = False
- with torch.no_grad():
- for stage, task in self.exchanger.get_cpu_tasks(self.cpu_idx):
- if had_error:
- continue # pragma: no cover
- try:
- if stage == 0:
- gpu_idx = None
- batch_id = next_batch_id
- debug("preprocess start for", batch_id)
- next_batch_id += 1
- docs = task
- else:
- gpu_idx, batch_id, result = task
- debug("postprocess start for", batch_id)
- docs = active_batches.pop(batch_id)
- gpu_pipe = stages[stage - 1]["gpu_component"]
- docs = gpu_pipe.postprocess(docs, result) # type: ignore
-
- for component in stages[stage]["cpu_components"]:
- if hasattr(component, "batch_process"):
- docs = component.batch_process(docs)
- else:
- docs = [component(doc) for doc in docs]
-
- gpu_pipe = stages[stage]["gpu_component"]
- if gpu_pipe is not None:
- preprocessed = gpu_pipe.make_batch(docs) # type: ignore
- active_batches[batch_id] = docs
- if gpu_idx is None:
- gpu_idx = batch_id % len(self.exchanger.gpu_worker_devices)
- collated = gpu_pipe.collate( # type: ignore
- preprocessed,
- device=self.exchanger.gpu_worker_devices[gpu_idx],
- )
- self.exchanger.put_gpu(
- item=(self.cpu_idx, batch_id, collated),
- idx=gpu_idx,
- stage=stage,
- )
- batch_id += 1
- debug("preprocess end for", batch_id)
- else:
- self.exchanger.put_results((docs, self.cpu_idx, gpu_idx))
- debug("postprocess end for", batch_id)
- except BaseException as e:
- had_error = True
- import traceback
-
- print(traceback.format_exc(), flush=True)
- self.exchanger.put_results((e, self.cpu_idx, None))
- # We need to drain the queues of GPUWorker fed inputs (pre-moved to GPU)
- # to ensure no tensor allocated on producer processes (CPUWorker via
- # collate) are left in consumer processes
- debug("Start draining CPU worker", self.cpu_idx)
- [None for _ in self.exchanger.get_cpu_tasks(self.cpu_idx)]
- debug(f"CPU worker {self.cpu_idx} is about to stop")
-
- def run(self):
- self._run()
- self.model = None
- gc.collect()
- torch.cuda.empty_cache()
-
-
-class GPUWorker(mp.Process):
- def __init__(
- self,
- gpu_idx,
- exchanger: Exchanger,
- gpu_pipe_names: List[str],
- model: Any,
- device: Union[str, torch.device],
- ):
- super().__init__()
-
- self.device = device
- self.gpu_idx = gpu_idx
- self.exchanger = exchanger
-
- self.gpu_pipe_names = gpu_pipe_names
- self.model = model
- self.device = device
-
- def _run(self):
- debug("GPU worker", self.gpu_idx, "started")
- mp._prctl_pr_set_pdeathsig(signal.SIGINT)
- had_error = False
-
- model = self.model.to(self.device)
- stage_components = [model.get_pipe(name) for name in self.gpu_pipe_names]
- del model
- with torch.no_grad():
- for stage, task in self.exchanger.get_gpu_tasks(self.gpu_idx):
- if had_error:
- continue # pragma: no cover
- try:
- cpu_idx, batch_id, batch = task
- debug("forward start for", batch_id)
- component = stage_components[stage]
- res = component.module_forward(batch)
- del batch, task
- # TODO set non_blocking=True here
- res = {
- key: val.to("cpu") if not isinstance(val, int) else val
- for key, val in res.items()
- }
- self.exchanger.put_cpu(
- item=(self.gpu_idx, batch_id, res),
- stage=stage + 1,
- idx=cpu_idx,
- )
- debug("forward end for", batch_id)
- except BaseException as e:
- had_error = True
- self.exchanger.put_results((e, None, self.gpu_idx))
- import traceback
-
- print(traceback.format_exc(), flush=True)
- task = batch = res = None # noqa
- # We need to drain the queues of CPUWorker fed inputs (pre-moved to GPU)
- # to ensure no tensor allocated on producer processes (CPUWorker via
- # collate) are left in consumer processes
- debug("Start draining GPU worker", self.gpu_idx)
- [None for _ in self.exchanger.get_gpu_tasks(self.gpu_idx)]
- debug(f"GPU worker {self.gpu_idx} is about to stop")
-
- def run(self):
- self._run()
- self.model = None
- gc.collect()
- torch.cuda.empty_cache()
-
-
-DEFAULT_MAX_CPU_WORKERS = 4
+from .base import Accelerator
@registry.accelerator.register("multiprocessing")
class MultiprocessingAccelerator(Accelerator):
"""
- If you have multiple CPU cores, and optionally multiple GPUs, we provide a
- `multiprocessing` accelerator that allows to run the inference on multiple
- processes.
-
- This accelerator dispatches the batches between multiple workers
- (data-parallelism), and distribute the computation of a given batch on one or two
- workers (model-parallelism). This is done by creating two types of workers:
-
- - a `CPUWorker` which handles the non deep-learning components and the
- preprocessing, collating and postprocessing of deep-learning components
- - a `GPUWorker` which handles the forward call of the deep-learning components
-
- The advantage of dedicating a worker to the deep-learning components is that it
- allows to prepare multiple batches in parallel in multiple `CPUWorker`, and ensure
- that the `GPUWorker` never wait for a batch to be ready.
-
- The overall architecture described in the following figure, for 3 CPU workers and 2
- GPU workers.
-
-
- 
- 
+ 
+