flow_inference.inference.Inference

class flow_inference.inference.Inference(download_repo_name, hf_token, trocr_model='microsoft/trocr-small-handwritten', target_image_size=None, stop_on_fail=False, splits=None, push_to_hub=True, private_repo=True, upload_repo_name=None, upload_mode='new_repo', allow_source_repo_update=False)[source]

Coordinate dataset loading, TrOCR inference, result writeback, and upload.

The pipeline loads records from a Hugging Face dataset, runs line-level inference with a configured TrOCR model, writes predictions back into the corresponding DataFrames, and can publish the updated dataset to a target Hugging Face repository.

Parameters:

  • download_repo_name (str)

  • hf_token (str | None)

  • target_image_size (Tuple[int, int])

  • stop_on_fail (bool)

  • splits (List[str] | None)

  • push_to_hub (bool)

  • private_repo (bool)

  • upload_repo_name (str | None)

  • upload_mode (Literal['new_repo', 'replace', 'update'])

  • allow_source_repo_update (bool)

__init__(download_repo_name, hf_token, trocr_model='microsoft/trocr-small-handwritten', target_image_size=None, stop_on_fail=False, splits=None, push_to_hub=True, private_repo=True, upload_repo_name=None, upload_mode='new_repo', allow_source_repo_update=False)[source]

Initialize the inference pipeline.

Parameters:

  • download_repo_name (str) – Hugging Face dataset repository to download.

  • hf_token (Optional[str]) – Optional Hugging Face token used for private repositories.

  • trocr_model – Hugging Face model ID or local path for the TrOCR model.

  • target_image_size (Tuple[int, int]) – Optional target image size as (width, height).

  • stop_on_fail (bool) – Whether to stop the workflow after an inference failure.

  • splits (Optional[List[str]]) – Dataset splits to process. Defaults to ["train"].

  • push_to_hub (bool) – Whether to upload the updated dataset after inference.

  • private_repo (bool) – Whether to create the upload repository as private.

  • upload_repo_name (Optional[str]) – Optional target dataset repository. Defaults to the source repository.

  • upload_mode (Literal['new_repo', 'replace', 'update']) – Upload behavior. Use "new_repo", "replace", or "update".

  • allow_source_repo_update (bool) – Whether uploading back into the source repository is allowed.

Raises:

RuntimeError – If the TrOCR model or processor cannot be loaded.

Return type:

None

Methods

__init__(download_repo_name, hf_token[, ...])

Initialize the inference pipeline.

perform_inference()

Run the complete inference workflow.

run_inference(records, model, processor, device)

Run TrOCR inference on selected image records.

save_results(inferred_lines, original_df)

Save inference predictions to a DataFrame.

write_inference_to_dataframe(inferred_lines, ...)

Write inference predictions into a timestamped DataFrame column.