emmi_inference.api
==================

.. py:module:: emmi_inference.api


Functions
---------

.. autoapisummary::

   emmi_inference.api.load_yaml
   emmi_inference.api.run_inference_from_config
   emmi_inference.api.setup_environment
   emmi_inference.api.load_components
   emmi_inference.api.create_writer
   emmi_inference.api.process_dataset_sample_by_sample
   emmi_inference.api.run_post_upload


Module Contents
---------------

.. py:function:: load_yaml(path)

   Reads a YAML file and returns a dictionary.

   :param path: The path to the YAML file.

   :returns:

             - Dictionary mapping keys to values.


.. py:function:: run_inference_from_config(config_path, model_type, model_path, destination, device = DeviceEnum.AUTO, autocast = False, dtype = None, upload_mode = UploadMode.SPOOL, remove_local_spool = False, upload_batch_size = 100, upload_workers = 8)

   Run inference based on the input config and the provided arguments.

   :param config_path: Path to config file (YAML) to use.
   :param model_type: A string name of the model type to use.
   :param model_path: Path to the model checkpoint.
   :param device: Device to use. Defaults to "auto".
   :param autocast: Either or not autocast to be enabled. Works only for CUDA. Defaults to False.
   :param dtype: Datatype to use. Defaults to None.
   :param destination: Output destination. Can be a local path (e.g., './results') or a remote URI
                       (e.g., 's3://bucket/key'). If not set, uses 'save_dir' from the config file for local-only output.
   :param upload_mode: Upload mode when `destination` is set. Defaults to 'spool'.
   :param remove_local_spool: Remove local spool/cache after successful upload.
   :param upload_batch_size: Batch size for 'spool' mode uploads.
   :param upload_workers: Number of parallel upload workers.


.. py:function:: setup_environment(device, autocast, dtype)

   Resolves device, dtype and autocast settings.

   :param device: Device to use.
   :param autocast: Either or not use autocast.
   :param dtype: Datatype to use. Defaults to None.

   :returns: Device string, dtype string or None, and autocast settings.
   :rtype: tuple


.. py:function:: load_components(config, model_type, model_path, device_str, torch_dtype, effective_autocast)

   Instantiates and returns the core components like Runner, Dataset, etc.

   :param config: Inout configuration as a dictionary.
   :param model_type: A string name of the model type to use.
   :param model_path: Path to the model checkpoint.
   :param device_str: Device string to use, e.g. "cuda".
   :param torch_dtype: Datatype to use. Defaults to None.
   :param effective_autocast: Either or not use autocast.

   :returns: an inference runner and a dataset.
   :rtype: tuple


.. py:function:: create_writer(config_params, cli_destination, local_spool_dir, upload_mode, remove_local_spool, upload_workers, upload_batch_size)

   Determines the final destination and instantiates the correct
   writer (AsyncWriter or S3SpoolingUploader).

   :returns: (writer_context, final_destination, is_remote_destination)


.. py:function:: process_dataset_sample_by_sample(*, runner, dataset, writer, batch_simplification, batch_keys_to_drop, batch_keys_to_keep)

   Executes inference over the input dataset.
   Even though this method doesn't raise exceptions directly it may log errors and raise uncaught exceptions.

   :param runner: Inference runner to use.
   :param dataset: Dataset to be processed.
   :param writer: Results writer to use.
   :param batch_simplification: Dictionary of batch simplifications.
   :param batch_keys_to_drop: A set of keys to drop from the batch.
   :param batch_keys_to_keep: A set of keys to keep in the batch.


.. py:function:: run_post_upload(is_remote, upload_mode, local_spool_dir, final_destination, upload_workers, remove_local_spool)

   Uploads local data from the `local_spool_dir` to `final_destination` via S3FileUploader.
   This method is executed _after_ the inference is done and is valid for S3 uploads only.

   :param is_remote: Either the target is remote or local. If False - there will be no upload.
   :param upload_mode: Upload mode to use. If anything but "after" - there will be no upload.
   :param local_spool_dir: Local Spooling directory.
   :param final_destination: Desired final destination.
   :param upload_workers: Number of worker processes to use.
   :param remove_local_spool: Whether to remove the local spool directory after upload.