1use anyhow::Result;
2use futures::stream::{self, StreamExt};
3use indexmap::IndexMap;
4use kdam::{BarExt, tqdm};
5
6use crate::{Example, Prediction, core::MetaSignature};
7
8#[allow(async_fn_in_trait)]
9pub trait Module: Send + Sync {
10 async fn forward(&self, inputs: Example) -> Result<Prediction>;
11
12 async fn batch(
13 &self,
14 inputs: Vec<Example>,
15 max_concurrency: usize,
16 display_progress: bool,
17 ) -> Result<Vec<Prediction>> {
18 let total = inputs.len();
19 let mut pb = if display_progress {
20 Some(tqdm!(total = total, desc = "Processing"))
21 } else {
22 None
23 };
24
25 let indexed_results: Vec<(usize, Result<Prediction>)> =
27 stream::iter(inputs.into_iter().enumerate())
28 .map(|(idx, example)| async move {
29 let result = self.forward(example).await;
30 (idx, result)
31 })
32 .buffer_unordered(max_concurrency)
33 .inspect(|_| {
34 if let Some(ref mut progress) = pb {
35 let _ = progress.update(1);
36 }
37 })
38 .collect()
39 .await;
40
41 let mut indexed_results = indexed_results;
43 indexed_results.sort_by_key(|(idx, _)| *idx);
44
45 let mut predictions = Vec::with_capacity(total);
47 for (_, result) in indexed_results {
48 predictions.push(result?);
49 }
50
51 Ok(predictions)
52 }
53}
54
55#[allow(unused_variables)]
56pub trait Optimizable {
57 fn get_signature(&self) -> &dyn MetaSignature {
58 todo!()
59 }
60
61 fn parameters(&mut self) -> IndexMap<String, &mut dyn Optimizable>;
62
63 fn update_signature_instruction(&mut self, instruction: String) -> anyhow::Result<()> {
64 todo!()
65 }
66}