burn_lm_registry/
lib.rs

1use burn_lm_inference::*;
2use burn_lm_macros::inference_server_registry;
3use std::{collections::HashMap, sync::Arc};
4
5pub type Channel<B> = MutexChannel<B>;
6
7pub type DynClients = HashMap<&'static str, Box<dyn InferencePlugin>>;
8
9// Register model crates
10#[inference_server_registry(
11    server(
12        crate_namespace = "burn_lm_llama::server::llama3",
13        server_type = "Llama3InstructServer<InferenceBackend>",
14    ),
15    server(
16        crate_namespace = "burn_lm_llama::server::llama3",
17        server_type = "Llama31InstructServer<InferenceBackend>",
18    ),
19    server(
20        crate_namespace = "burn_lm_llama::server::llama3",
21        server_type = "Llama321bInstructServer<InferenceBackend>",
22    ),
23    server(
24        crate_namespace = "burn_lm_llama::server::llama3",
25        server_type = "Llama323bInstructServer<InferenceBackend>",
26    ),
27    server(
28        crate_namespace = "burn_lm_llama::server::tiny",
29        server_type = "TinyLlamaServer<InferenceBackend>",
30    ),
31    server(
32        crate_namespace = "burn_lm_parrot",
33        server_type = "ParrotServer<InferenceBackend>",
34    )
35)]
36#[derive(Debug)]
37pub struct Registry {
38    clients: Arc<DynClients>,
39}