cortenforge_tools/warehouse_commands/
common.rs

1use std::{borrow::Cow, fmt};
2
3#[allow(dead_code)]
4#[derive(Clone, Copy, Debug)]
5pub enum WarehouseStore {
6    Memory,
7    Mmap,
8    Stream,
9}
10
11impl WarehouseStore {
12    pub fn as_str(&self) -> &'static str {
13        match self {
14            WarehouseStore::Memory => "memory",
15            WarehouseStore::Mmap => "mmap",
16            WarehouseStore::Stream => "stream",
17        }
18    }
19}
20
21impl fmt::Display for WarehouseStore {
22    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
23        f.write_str(self.as_str())
24    }
25}
26
27#[allow(dead_code)]
28#[derive(Clone, Copy, Debug)]
29pub enum ModelKind {
30    Tiny,
31    Big,
32}
33
34impl ModelKind {
35    pub fn as_str(&self) -> &'static str {
36        match self {
37            ModelKind::Tiny => "tiny",
38            ModelKind::Big => "big",
39        }
40    }
41}
42
43impl fmt::Display for ModelKind {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.write_str(self.as_str())
46    }
47}
48
49#[allow(dead_code)]
50#[derive(Clone, Debug)]
51pub struct CmdConfig<'a> {
52    pub manifest: Cow<'a, str>,
53    pub store: WarehouseStore,
54    pub prefetch: Option<usize>, // used only when store == Stream
55    pub model: ModelKind,
56    pub batch_size: usize,
57    pub log_every: usize,
58    pub wgpu_backend: Cow<'a, str>,
59    pub wgpu_adapter: Option<Cow<'a, str>>,
60    pub extra_args: Cow<'a, str>,
61}
62
63impl<'a> CmdConfig<'a> {
64    pub fn with_manifest<T: Into<Cow<'a, str>>>(mut self, manifest: T) -> Self {
65        self.manifest = manifest.into();
66        self
67    }
68
69    pub fn with_store(mut self, store: WarehouseStore) -> Self {
70        self.store = store;
71        self
72    }
73
74    pub fn with_prefetch(mut self, prefetch: Option<usize>) -> Self {
75        self.prefetch = prefetch;
76        self
77    }
78
79    pub fn with_model(mut self, model: ModelKind) -> Self {
80        self.model = model;
81        self
82    }
83
84    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
85        self.batch_size = batch_size;
86        self
87    }
88
89    pub fn with_log_every(mut self, log_every: usize) -> Self {
90        self.log_every = log_every;
91        self
92    }
93
94    pub fn with_backend<T: Into<Cow<'a, str>>>(mut self, backend: T) -> Self {
95        self.wgpu_backend = backend.into();
96        self
97    }
98
99    pub fn with_adapter<T: Into<Cow<'a, str>>>(mut self, adapter: T) -> Self {
100        self.wgpu_adapter = Some(adapter.into());
101        self
102    }
103
104    pub fn with_extra_args<T: Into<Cow<'a, str>>>(mut self, extra_args: T) -> Self {
105        self.extra_args = extra_args.into();
106        self
107    }
108}
109
110#[allow(dead_code)]
111pub const DEFAULT_CONFIG: CmdConfig<'static> = CmdConfig {
112    manifest: Cow::Borrowed("artifacts/tensor_warehouse/v<version>/manifest.json"),
113    store: WarehouseStore::Stream,
114    prefetch: Some(8),
115    model: ModelKind::Big,
116    batch_size: 32,
117    log_every: 1,
118    wgpu_backend: Cow::Borrowed("vulkan"),
119    wgpu_adapter: None,
120    extra_args: Cow::Borrowed(""),
121};
122
123impl Default for CmdConfig<'_> {
124    fn default() -> Self {
125        DEFAULT_CONFIG
126    }
127}
cortenforge_tools/warehouse_commands/common.rs

cortenforge_tools/warehouse_commands/
common.rs
