Skip to main content

surrealml_core/storage/header/
engine.rs

1//! Defines the placeholder for the type of model engine in the header.
2use std::fmt;
3
4/// Defines the type of engine being used to run the model.
5///
6/// # Fields
7/// * `Native` - The native engine which will be native rust and linfa.
8/// * `PyTorch` - The PyTorch engine which will be PyTorch and tch-rs.
9/// * `Undefined` - The undefined engine which will be used when the engine is not defined.
10#[derive(Debug, PartialEq)]
11pub enum Engine {
12	Native,
13	PyTorch,
14	Undefined,
15}
16
17impl Engine {
18	/// Creates a new `Engine` struct with the undefined engine.
19	///
20	/// # Returns
21	/// A new `Engine` struct with the undefined engine.
22	pub fn fresh() -> Self {
23		Engine::Undefined
24	}
25
26	/// Creates a new `Engine` struct from a string.
27	///
28	/// # Arguments
29	/// * `engine` - The engine as a string.
30	///
31	/// # Returns
32	/// A new `Engine` struct.
33	pub fn from_string(engine: String) -> Self {
34		match engine.as_str() {
35			"native" => Engine::Native,
36			"pytorch" => Engine::PyTorch,
37			_ => Engine::Undefined,
38		}
39	}
40}
41
42impl fmt::Display for Engine {
43	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44		match self {
45			Engine::Native => write!(f, "native"),
46			Engine::PyTorch => write!(f, "pytorch"),
47			Engine::Undefined => write!(f, ""),
48		}
49	}
50}