1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
use tract_core::internal::HashMap;
use std::{hash, cmp};
use pyo3::ToPyObject;
use tract_core::tensor::Tensor;
use std::sync::Arc;
use smallvec::SmallVec;

pub enum Types {
    PyModel,
    TFModel
}

#[derive(Debug, PartialEq, PartialOrd, Clone)]
pub struct ORequest<T>
{
    pub body: T
}

impl<T> ORequest<T> {
    pub fn with_body(body: T) -> Self {
        ORequest { body }
    }
}

#[derive(Debug, PartialEq, PartialOrd, Clone)]
pub struct OResponse<T>
{
    pub body: T
}

impl<T> OResponse<T> {
    pub fn with_body(body: T) -> Self {
        OResponse { body }
    }
}

pub(crate) trait ORequestBase<T> {}
pub(crate) trait OResponseBase<T> {}

impl<T> ORequestBase<T> for T {}
impl<T> OResponseBase<T> for T {}

#[derive(Default, Debug, Clone)]
pub struct PyModelRequest<K, V, T>
    where K: hash::Hash + cmp::Eq + Default + ToPyObject,
          V: Default + ToPyObject,
          T: Default + ToPyObject {
    pub args: HashMap<K, V>,
    pub kwargs: HashMap<&'static str, T>
}

impl<K, V, T> PyModelRequest<K, V, T>
    where K: hash::Hash + cmp::Eq + Default + ToPyObject,
          V: Default + ToPyObject,
          T: Default + ToPyObject {
    pub fn new() -> Self {
        PyModelRequest { ..Default::default() }
    }

    pub fn with_args(mut self, args: HashMap<K, V>) -> Self {
        self.args = args;
        self
    }

    pub fn with_kwargs(mut self, kwargs: HashMap<&'static str, T>) -> Self {
        self.kwargs = kwargs;
        self
    }
}

//#[derive(Default, Debug, Clone)]
//pub struct PyModelResponse {
//    response: PyDict,
//}
//
//impl PyModelResponse {
//    pub fn new() -> Self { PyModelResponse { ..Default::default() } }
//}

#[derive(Default, Debug)]
pub struct TFRequest {
    pub input: Tensor
}

impl TFRequest {
    pub fn new() ->  Self { TFRequest { .. Default::default()} }
}

#[derive(Default, Debug)]
pub struct TFResponse {
    pub output: SmallVec<[Arc<Tensor>; 4]>
}

impl TFResponse {
    pub fn new() ->  Self { TFResponse  { .. Default::default()} }

    pub fn with_output(mut self, output: SmallVec<[Arc<Tensor>; 4]>) -> Self {
        self.output = output;
        self
    }
}