burn_lm_inference/

job.rs

use std::{
    any::Any,
    io::Write,
    marker::PhantomData,
    sync::{
        atomic::{AtomicBool, Ordering},
        mpsc::SyncSender,
        Arc,
    },
};

use crate::{Message, Prompt};

/// Defines a job to be run during inference.
pub struct InferenceJob {
    /// The task to be performed by the job.
    pub task: InferenceTask,
    /// The emitter for the current job.
    pub emitter: GeneratedItemEmitter,
}

/// An emitter is responsible to send [generated items](GeneratedItem) to the [inference job](InferenceJob)
/// channel.
pub struct GeneratedItemEmitter {
    sender: SyncSender<Msg>,
    done: Arc<AtomicBool>,
}

/// The potential tasks that can be executed by an [inference job](InferenceJob) using the
/// [inference server](crate::InferenceServer).
pub enum InferenceTask {
    /// A single message to be processed by the server.
    Message(Message),
    /// Multiple messages to be processed by the server.
    ///
    /// This could be useful to restore a previous session based on text history.
    Context(Vec<Message>),
    /// Run with a simple prompt.
    Prompt(Prompt),
}

/// Defines all the potential items that can be generated by an [inference job](InferenceJob).
///
/// For now there is only text generation, but this could be extented to other kind of output
/// artifacts.
pub enum GeneratedItem {
    /// Generated text includes intermediary tokens and doesn't mark the end of a text generation
    /// job.
    Text(String),
}

impl InferenceJob {
    /// Start a new inference job and process it on another thread.
    ///
    /// When a task is performed for the current job, it should be registered using the
    /// [completed method](Self::completed).
    pub fn create<L: InferenceJobListener>(
        task: InferenceTask,
        listener: L,
    ) -> (Self, JobHandle<L>) {
        let (emitter, handle) = GeneratedItemEmitter::init(listener);

        (Self { task, emitter }, handle)
    }
}

impl GeneratedItemEmitter {
    pub fn init<L: InferenceJobListener>(mut listener: L) -> (Self, JobHandle<L>) {
        let (sender, receiver) = std::sync::mpsc::sync_channel::<Msg>(1);

        let handle = JobHandle {
            sender: sender.clone(),
            _c: PhantomData,
        };
        let done = Arc::new(AtomicBool::new(false));
        let emitter = GeneratedItemEmitter {
            sender,
            done: done.clone(),
        };

        // TODO: We could use a threadpool for inference jobs.
        std::thread::spawn(move || {
            for msg in receiver {
                match msg {
                    Msg::Text(text) => listener.on_text(text),
                    Msg::Finished(c) => {
                        let result = listener.on_finished();
                        let result: Box<dyn Any + Send> = Box::new(result);
                        c.send(result).unwrap();
                        done.store(true, Ordering::Relaxed);
                        return;
                    }
                }
            }
        });

        (emitter, handle)
    }

    /// Register the completion of an [inference generation](InferenceGeneration).
    pub fn completed(&self, item: GeneratedItem) {
        if !self.done.load(Ordering::Relaxed) {
            let msg = match item {
                GeneratedItem::Text(text) => Msg::Text(text),
            };
            self.sender.send(msg).unwrap();
        }
    }
}

/// An inference job listener receive events while the [inference job](InferenceJob) is running.
pub trait InferenceJobListener: Send + 'static {
    /// The item that is returned by the listener.
    type CompletedItem: Send;

    /// Called when new text is generated from an [inference job](InferenceJob).
    fn on_text(&mut self, text: String);

    /// Called when the job is finished.
    ///
    /// The inference job listener can return an item when a job is finished.
    /// The item is going to be available through the [job handle finished method](JobHandle::finished).
    fn on_finished(self) -> Self::CompletedItem;
}

#[derive(Default)]
/// The text generation listener accumulate the generated text in a string that can be
/// obtained at the end of the job with the [handle finished method](JobHandle::finished).
pub struct TextGenerationListener {
    value: String,
}

impl InferenceJobListener for TextGenerationListener {
    type CompletedItem = String;

    fn on_text(&mut self, text: String) {
        self.value += &text;
    }

    fn on_finished(self) -> Self::CompletedItem {
        self.value
    }
}

#[derive(Default)]
/// The stdout listener directly writes the intermediary [generated item](GeneratedItem) to
/// [std::io::stdout].
pub struct StdOutListener {}

impl InferenceJobListener for StdOutListener {
    type CompletedItem = ();

    fn on_text(&mut self, text: String) {
        let mut io = std::io::stdout();

        write!(io, "{text}").unwrap();
        io.flush().unwrap();
    }

    fn on_finished(self) -> Self::CompletedItem {}
}

/// The handle returned by [InferenceJob::start].
///
/// This handle should be used to indicate when a job is finished using the
/// [join method](JobHandle::join).
pub struct JobHandle<C: InferenceJobListener> {
    sender: SyncSender<Msg>,
    _c: PhantomData<C>,
}

impl<C: InferenceJobListener> JobHandle<C> {
    /// Wait for the job to complete and returns the
    /// [InferenceJobListener::CompletedItem] from the job listener.
    ///
    /// # Warning
    ///
    /// Make sure to actually launch the inference job before calling this method, otherwise this
    /// method will never complete.
    pub fn join(&self) -> C::CompletedItem {
        let (sender, rec) = std::sync::mpsc::sync_channel(1);
        self.sender.send(Msg::Finished(sender)).unwrap();

        match rec.recv() {
            Ok(any) => *any.downcast().unwrap(),
            Err(err) => panic!("{err}"),
        }
    }
}

enum Msg {
    Text(String),
    Finished(SyncSender<Box<dyn Any + Send>>),
}