tract_core/
framework.rs

1//! Enforce consistent API between the implemented Frameworks importers.
2use crate::internal::*;
3use std::fmt::Debug;
4use std::io::Read;
5use std::path::Path;
6
7/// A Framework that translate its model to tract core model.
8///
9/// The ProtoModel is the parsed representation of the imported model. It does
10/// not have to be Protobuf based.
11pub trait Framework<ProtoModel, Model>: Send + Sync
12where
13    ProtoModel: Debug,
14    Model: Default,
15{
16    /// Parse a proto model from a reader.
17    fn proto_model_for_read(&self, reader: &mut dyn Read) -> TractResult<ProtoModel>;
18
19    /// Translate a proto model into a model.
20    fn model_for_proto_model(&self, proto: &ProtoModel) -> TractResult<Model> {
21        self.model_for_proto_model_with_model_template(proto, Model::default())
22    }
23
24    /// Translate a proto model into a model, with some symbols already listed.
25    fn model_for_proto_model_with_model_template(
26        &self,
27        proto: &ProtoModel,
28        template: Model,
29    ) -> TractResult<Model>;
30
31    /// Read a proto model from a filename.
32    fn proto_model_for_path(&self, p: impl AsRef<Path>) -> TractResult<ProtoModel> {
33        let mut r = std::fs::File::open(p.as_ref())
34            .with_context(|| format!("Could not open {:?}", p.as_ref()))?;
35        self.proto_model_for_read(&mut r)
36    }
37
38    /// Read a model from a reader
39    fn model_for_read(&self, r: &mut dyn Read) -> TractResult<Model> {
40        let proto_model = self.proto_model_for_read(r).context("Reading proto model")?;
41        self.model_for_proto_model(&proto_model).context("Translating proto model to model")
42    }
43
44    /// Build a model from a filename.
45    fn model_for_path(&self, p: impl AsRef<Path>) -> TractResult<Model> {
46        let mut r = std::fs::File::open(p.as_ref())
47            .with_context(|| format!("Could not open {:?}", p.as_ref()))?;
48        self.model_for_read(&mut r)
49    }
50}