docs: more rust documentation

Author: decahedron1

build.rs

@@ -530,7 +530,7 @@ fn generate_bindings(include_dir: &Path) {
 }
 
 fn main() {
-	if !std::env::var("DOCS_RS").is_ok() {
+	if std::env::var("DOCS_RS").is_err() {
 		let (install_dir, needs_link) = prepare_libort_dir();
 
 		let include_dir = install_dir.join("include");

src/download.rs

@@ -1,6 +1,8 @@
 pub mod language;
 pub mod vision;
 
+/// Represents a type that returns an ONNX model URL.
 pub trait ModelUrl {
+	/// Returns the model URL associated with this model.
 	fn fetch_url(&self) -> &'static str;
 }

src/download/language.rs

@@ -1,3 +1,5 @@
+//! Models for language understanding.
+
 pub mod machine_comprehension;
 
-pub use machine_comprehension::MachineComprehension;
+pub use machine_comprehension::{MachineComprehension, RoBERTa, GPT2};

src/download/language/machine_comprehension.rs

@@ -1,5 +1,7 @@
 #![allow(clippy::upper_case_acronyms)]
 
+//! Models for machine language comprehension.
+
 use crate::download::ModelUrl;
 
 /// Machine comprehension models.
@@ -20,14 +22,18 @@ pub enum MachineComprehension {
 /// Large transformer-based model that predicts sentiment based on given input text.
 #[derive(Debug, Clone)]
 pub enum RoBERTa {
+	/// Base RoBERTa model.
 	RoBERTaBase,
+	/// RoBERTa model for sequence classification.
 	RoBERTaSequenceClassification
 }
 
 /// Generates synthetic text samples in response to the model being primed with an arbitrary input.
 #[derive(Debug, Clone)]
 pub enum GPT2 {
+	/// Base GPT-2 model.
 	GPT2,
+	/// GPT-2 model with a causal LM head.
 	GPT2LmHead
 }
 

src/download/vision.rs

@@ -1,3 +1,5 @@
+//! Models for computer vision.
+
 pub mod body_face_gesture_analysis;
 pub mod domain_based_image_classification;
 pub mod image_classification;
@@ -6,6 +8,6 @@ pub mod object_detection_image_segmentation;
 
 pub use body_face_gesture_analysis::BodyFaceGestureAnalysis;
 pub use domain_based_image_classification::DomainBasedImageClassification;
-pub use image_classification::ImageClassification;
-pub use image_manipulation::ImageManipulation;
+pub use image_classification::{ImageClassification, InceptionVersion, ResNet, ResNetV1, ResNetV2, ShuffleNetVersion, Vgg};
+pub use image_manipulation::{FastNeuralStyleTransferStyle, ImageManipulation};
 pub use object_detection_image_segmentation::ObjectDetectionImageSegmentation;

src/download/vision/body_face_gesture_analysis.rs

@@ -1,5 +1,8 @@
+//! Models for body, face, & gesture analysis.
+
 use crate::download::ModelUrl;
 
+/// Models for body, face, & gesture analysis.
 #[derive(Debug, Clone)]
 pub enum BodyFaceGestureAnalysis {
 	/// A CNN based model for face recognition which learns discriminative features of faces and produces embeddings for

src/download/vision/domain_based_image_classification.rs

@@ -1,5 +1,8 @@
+//! Models for domain-based image classification.
+
 use crate::download::ModelUrl;
 
+/// Models for domain-based image classification.
 #[derive(Debug, Clone)]
 pub enum DomainBasedImageClassification {
 	/// Handwritten digit prediction using CNN.

src/download/vision/image_classification.rs

@@ -1,7 +1,10 @@
+//! Models for image classification.
+
 #![allow(clippy::upper_case_acronyms)]
 
 use crate::download::ModelUrl;
 
+/// Models for image classification.
 #[derive(Debug, Clone)]
 pub enum ImageClassification {
 	/// Image classification aimed for mobile targets.
@@ -73,10 +76,15 @@ pub enum ResNet {
 
 #[derive(Debug, Clone)]
 pub enum ResNetV1 {
+	/// ResNet v1 with 18 layers.
 	ResNet18,
+	/// ResNet v1 with 34 layers.
 	ResNet34,
+	/// ResNet v1 with 50 layers.
 	ResNet50,
+	/// ResNet v1 with 101 layers.
 	ResNet101,
+	/// ResNet v1 with 152 layers.
 	ResNet152
 }
 
@@ -109,6 +117,7 @@ pub enum Vgg {
 /// power.
 #[derive(Debug, Clone)]
 pub enum ShuffleNetVersion {
+	/// The original ShuffleNet.
 	V1,
 	/// ShuffleNetV2 is an improved architecture that is the state-of-the-art in terms of speed and accuracy tradeoff
 	/// used for image classification.

src/error.rs

@@ -1,3 +1,5 @@
+//! Types and helpers for handling ORT errors.
+
 use std::{io, path::PathBuf, string};
 
 use thiserror::Error;
@@ -7,9 +9,11 @@ use super::{char_p_to_string, ort, sys, tensor::TensorElementDataType};
 /// Type alias for the Result type returned by ORT functions.
 pub type OrtResult<T> = std::result::Result<T, OrtError>;
 
+/// An enum of all errors returned by ORT functions.
 #[non_exhaustive]
 #[derive(Error, Debug)]
 pub enum OrtError {
+	/// An error occurred when converting an FFI C string to a Rust `String`.
 	#[error("Failed to construct Rust String")]
 	FfiStringConversion(OrtApiError),
 	/// An error occurred while creating an ONNX environment.

src/execution_providers.rs

@@ -16,6 +16,10 @@ extern "C" {
 	pub(crate) fn OrtSessionOptionsAppendExecutionProvider_DML(options: *mut sys::OrtSessionOptions, device_id: std::os::raw::c_int) -> sys::OrtStatusPtr;
 }
 
+/// Execution provider container. See [the ONNX Runtime docs](https://onnxruntime.ai/docs/execution-providers/) for more
+/// info on execution providers. Execution providers are actually registered via the `with_execution_providers()`
+/// functions [`crate::SessionBuilder`] (per-session) or [`crate::EnvBuilder`] (default for all sessions in an
+/// environment).
 #[derive(Debug, Clone)]
 pub struct ExecutionProvider {
 	provider: String,
@@ -51,6 +55,9 @@ macro_rules! ep_options {
 }
 
 impl ExecutionProvider {
+	/// Creates an `ExecutionProvider` for the given execution provider name.
+	///
+	/// You probably want the dedicated methods instead, e.g. [`ExecutionProvider::cuda`].
 	pub fn new(provider: impl Into<String>) -> Self {
 		Self {
 			provider: provider.into(),
@@ -69,6 +76,8 @@ impl ExecutionProvider {
 		directml = "DmlExecutionProvider"
 	}
 
+	/// Returns `true` if this execution provider is available, `false` otherwise.
+	/// The CPU execution provider will always be available.
 	pub fn is_available(&self) -> bool {
 		let mut providers: *mut *mut c_char = std::ptr::null_mut();
 		let mut num_providers = 0;
@@ -90,9 +99,9 @@ impl ExecutionProvider {
 		false
 	}
 
-	/// Configure this execution provider with the given option.
-	pub fn with(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
-		self.options.insert(k.into(), v.into());
+	/// Configure this execution provider with the given option name and value
+	pub fn with(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
+		self.options.insert(name.into(), value.into());
 		self
 	}
 
@@ -198,7 +207,7 @@ pub(crate) fn apply_execution_providers(options: *mut sys::OrtSessionOptions, ex
 			"DmlExecutionProvider" => {
 				let device_id = init_args.get("device_id").map_or(0, |s| s.parse::<i32>().unwrap_or(0));
 				// TODO: extended options with OrtSessionOptionsAppendExecutionProviderEx_DML
-				let status = unsafe { OrtSessionOptionsAppendExecutionProvider_DML(options, device_id.into()) };
+				let status = unsafe { OrtSessionOptionsAppendExecutionProvider_DML(options, device_id) };
 				if status_to_result_and_log("DirectML", status).is_ok() {
 					return; // EP found
 				}

src/lib.rs

@@ -1,3 +1,5 @@
+#![doc = include_str!("../README.md")]
+
 pub mod download;
 pub mod environment;
 pub mod error;
@@ -42,6 +44,11 @@ lazy_static! {
 	};
 }
 
+/// Attempts to acquire the global OrtApi object.
+///
+/// # Panics
+///
+/// Panics if another thread panicked while holding the API lock, or if the ONNX Runtime API could not be initialized.
 pub fn ort() -> sys::OrtApi {
 	let mut api_ref = G_ORT_API.lock().expect("failed to acquire OrtApi lock; another thread panicked?");
 	let api_ref_mut: &mut *mut sys::OrtApi = api_ref.get_mut();
@@ -176,15 +183,20 @@ extern_system_fn! {
 	}
 }
 
-/// ONNX Runtime logging level.
+/// The minimum logging level. Logs will be handled by the `tracing` crate.
 #[derive(Debug)]
 #[cfg_attr(not(windows), repr(u32))]
 #[cfg_attr(windows, repr(i32))]
 pub enum LoggingLevel {
+	/// Verbose logging level. This will log *a lot* of messages!
 	Verbose = sys::OrtLoggingLevel_ORT_LOGGING_LEVEL_VERBOSE as OnnxEnumInt,
+	/// Info logging level.
 	Info = sys::OrtLoggingLevel_ORT_LOGGING_LEVEL_INFO as OnnxEnumInt,
+	/// Warning logging level. Recommended to receive potentially important warnings.
 	Warning = sys::OrtLoggingLevel_ORT_LOGGING_LEVEL_WARNING as OnnxEnumInt,
+	/// Error logging level.
 	Error = sys::OrtLoggingLevel_ORT_LOGGING_LEVEL_ERROR as OnnxEnumInt,
+	/// Fatal logging level.
 	Fatal = sys::OrtLoggingLevel_ORT_LOGGING_LEVEL_FATAL as OnnxEnumInt
 }
 
@@ -212,7 +224,7 @@ impl From<LoggingLevel> for sys::OrtLoggingLevel {
 /// The optimizations belonging to one level are performed after the optimizations of the previous level have been
 /// applied (e.g., extended optimizations are applied after basic optimizations have been applied).
 ///
-/// **All optimizations are enabled by default.**
+/// **All optimizations (i.e. [`GraphOptimizationLevel::Level3`]) are enabled by default.**
 ///
 /// # Online/offline mode
 /// All optimizations can be performed either online or offline. In online mode, when initializing an inference session,
@@ -233,6 +245,7 @@ impl From<LoggingLevel> for sys::OrtLoggingLevel {
 #[cfg_attr(not(windows), repr(u32))]
 #[cfg_attr(windows, repr(i32))]
 pub enum GraphOptimizationLevel {
+	/// Disables all graph optimizations.
 	Disable = sys::GraphOptimizationLevel_ORT_DISABLE_ALL as OnnxEnumInt,
 	/// Level 1 includes semantics-preserving graph rewrites which remove redundant nodes and redundant computation.
 	/// They run before graph partitioning and thus apply to all the execution providers. Available basic/level 1 graph
@@ -292,13 +305,13 @@ impl From<GraphOptimizationLevel> for sys::GraphOptimizationLevel {
 	}
 }
 
-/// Allocator type
+/// Execution provider allocator type.
 #[derive(Debug, Clone)]
 #[repr(i32)]
 pub enum AllocatorType {
-	/// Device allocator
+	/// Default device-specific allocator.
 	Device = sys::OrtAllocatorType_OrtDeviceAllocator,
-	/// Arena allocator
+	/// Arena allocator.
 	Arena = sys::OrtAllocatorType_OrtArenaAllocator
 }
 
@@ -311,17 +324,20 @@ impl From<AllocatorType> for sys::OrtAllocatorType {
 	}
 }
 
-/// Memory type
+/// Memory types for allocated memory.
 #[derive(Debug, Clone)]
 #[repr(i32)]
 pub enum MemType {
+	/// Any CPU memory used by non-CPU execution provider.
 	CPUInput = sys::OrtMemType_OrtMemTypeCPUInput,
+	/// CPU accessible memory outputted by non-CPU execution provider, i.e. CUDA_PINNED.
 	CPUOutput = sys::OrtMemType_OrtMemTypeCPUOutput,
-	/// Default memory type
+	/// The default allocator for an execution provider.
 	Default = sys::OrtMemType_OrtMemTypeDefault
 }
 
 impl MemType {
+	/// Temporary CPU accessible memory allocated by non-CPU execution provider, i.e. CUDA_PINNED.
 	pub const CPU: MemType = MemType::CPUOutput;
 }
 

src/metadata.rs

@@ -4,6 +4,7 @@ use std::{ffi::CString, os::raw::c_char};
 
 use super::{char_p_to_string, error::OrtResult, ortfree, ortsys, sys, OrtError};
 
+/// Container for model metadata, including name & producer information.
 pub struct Metadata {
 	metadata_ptr: *mut sys::OrtModelMetadata,
 	allocator_ptr: *mut sys::OrtAllocator
@@ -14,6 +15,7 @@ impl Metadata {
 		Metadata { metadata_ptr, allocator_ptr }
 	}
 
+	/// Gets the model description, returning an error if no description is present.
 	pub fn description(&self) -> OrtResult<String> {
 		let mut str_bytes: *mut c_char = std::ptr::null_mut();
 		ortsys![unsafe ModelMetadataGetDescription(self.metadata_ptr, self.allocator_ptr, &mut str_bytes) -> OrtError::GetModelMetadata; nonNull(str_bytes)];
@@ -23,6 +25,7 @@ impl Metadata {
 		Ok(value)
 	}
 
+	/// Gets the model producer name, returning an error if no producer name is present.
 	pub fn producer(&self) -> OrtResult<String> {
 		let mut str_bytes: *mut c_char = std::ptr::null_mut();
 		ortsys![unsafe ModelMetadataGetProducerName(self.metadata_ptr, self.allocator_ptr, &mut str_bytes) -> OrtError::GetModelMetadata; nonNull(str_bytes)];
@@ -32,6 +35,7 @@ impl Metadata {
 		Ok(value)
 	}
 
+	/// Gets the model name, returning an error if no name is present.
 	pub fn name(&self) -> OrtResult<String> {
 		let mut str_bytes: *mut c_char = std::ptr::null_mut();
 		ortsys![unsafe ModelMetadataGetGraphName(self.metadata_ptr, self.allocator_ptr, &mut str_bytes) -> OrtError::GetModelMetadata; nonNull(str_bytes)];
@@ -41,12 +45,14 @@ impl Metadata {
 		Ok(value)
 	}
 
+	/// Gets the model version, returning an error if no version is present.
 	pub fn version(&self) -> OrtResult<i64> {
 		let mut ver = 0i64;
 		ortsys![unsafe ModelMetadataGetVersion(self.metadata_ptr, &mut ver) -> OrtError::GetModelMetadata];
 		Ok(ver)
 	}
 
+	/// Fetch the value of a custom metadata key. Returns `Ok(None)` if the key is not found.
 	pub fn custom(&self, key: &str) -> OrtResult<Option<String>> {
 		let mut str_bytes: *mut c_char = std::ptr::null_mut();
 		let key_str = CString::new(key)?;

src/session.rs

@@ -1,3 +1,5 @@
+//! Contains the [`Session`] and [`SessionBuilder`] types for managing ONNX Runtime sessions and performing inference.
+
 #![allow(clippy::tabs_in_doc_comments)]
 
 #[cfg(not(target_family = "windows"))]
@@ -95,6 +97,7 @@ impl Drop for SessionBuilder {
 }
 
 impl SessionBuilder {
+	/// Creates a new session builder in the given environment.
 	pub fn new(env: &Arc<Environment>) -> OrtResult<Self> {
 		let mut session_options_ptr: *mut sys::OrtSessionOptions = std::ptr::null_mut();
 		ortsys![unsafe CreateSessionOptions(&mut session_options_ptr) -> OrtError::CreateSessionOptions; nonNull(session_options_ptr)];
@@ -673,14 +676,16 @@ impl Session {
 		Ok(())
 	}
 
+	/// Gets the session model metadata. See [`Metadata`] for more info.
 	pub fn metadata(&self) -> OrtResult<Metadata> {
 		let mut metadata_ptr: *mut sys::OrtModelMetadata = std::ptr::null_mut();
 		ortsys![unsafe SessionGetModelMetadata(self.session_ptr, &mut metadata_ptr) -> OrtError::GetModelMetadata; nonNull(metadata_ptr)];
 		Ok(Metadata::new(metadata_ptr, self.allocator_ptr))
 	}
 
-	/// Ends profiling for this session. Note that this must be explicitly called at the end of profiling, otherwise
-	/// the profiing file will be empty.
+	/// Ends profiling for this session.
+	///
+	/// Note that this must be explicitly called at the end of profiling, otherwise the profiing file will be empty.
 	#[cfg(feature = "profiling")]
 	pub fn end_profiling(&self) -> OrtResult<String> {
 		let mut profiling_name: *mut c_char = std::ptr::null_mut();

src/tensor.rs

@@ -148,6 +148,7 @@ impl_type_trait!(half::bf16, Bfloat16);
 /// would conflict with the implementations of [IntoTensorElementDataType] for primitive numeric
 /// types (which might implement [`AsRef<str>`] at some point in the future).
 pub trait Utf8Data {
+	/// Returns the contents of this value as a slice of UTF-8 bytes.
 	fn utf8_bytes(&self) -> &[u8];
 }