modelc 0.1.4

Compile model weight files to standalone executable binaries
Documentation
use std::convert::Infallible;
use std::sync::Arc;

use axum::{Json, extract::State, response::sse::{Event, Sse}};
use tokio_stream::wrappers::ReceiverStream;

use super::{AppState, ChatRequest, ChatResponse, CompleteRequest, CompleteResponse, InferRequest, InferResponse, Message, ModelInfo, StreamChunk};
use super::infer::{run_inference, run_text_inference};

pub(super) async fn infer(
    State(state): State<Arc<AppState>>,
    Json(req): Json<InferRequest>,
) -> Json<InferResponse> {
    if !req.inputs.is_empty() {
        let start = std::time::Instant::now();
        let outs: Vec<Vec<f32>> = req
            .inputs
            .iter()
            .map(|inp| run_inference(&state, inp))
            .collect();
        if state.profile {
            eprintln!("  batch infer: {} items in {:.3} ms", outs.len(), start.elapsed().as_secs_f64() * 1000.0);
        }
        return Json(InferResponse {
            output: None,
            outputs: Some(outs),
        });
    }

    let start = std::time::Instant::now();
    let output = run_inference(&state, &req.input);
    if state.profile {
        eprintln!("  infer: {:.3} ms", start.elapsed().as_secs_f64() * 1000.0);
    }
    Json(InferResponse {
        output: Some(output),
        outputs: None,
    })
}

pub(super) async fn model_info(State(state): State<Arc<AppState>>) -> Json<ModelInfo> {
    Json(ModelInfo {
        name: state.name.clone(),
        architecture: state.architecture.clone(),
        total_params: state.total_params,
        total_bytes: state.total_bytes,
        tensors: state.tensor_names.clone(),
    })
}

pub(super) async fn chat(
    State(state): State<Arc<AppState>>,
    Json(req): Json<ChatRequest>,
) -> Json<ChatResponse> {
    let prompt = req
        .messages
        .last()
        .map(|m| m.content.clone())
        .unwrap_or_default();
    let output = run_text_inference(&state, &prompt);
    Json(ChatResponse {
        message: Message {
            role: "assistant".to_string(),
            content: output,
        },
    })
}

pub(super) async fn complete(
    State(state): State<Arc<AppState>>,
    Json(req): Json<CompleteRequest>,
) -> Json<CompleteResponse> {
    let output = run_text_inference(&state, &req.prompt);
    Json(CompleteResponse { completion: output })
}

pub(super) async fn chat_stream(
    State(state): State<Arc<AppState>>,
    Json(req): Json<ChatRequest>,
) -> Sse<ReceiverStream<Result<Event, Infallible>>> {
    let prompt = req
        .messages
        .last()
        .map(|m| m.content.clone())
        .unwrap_or_default();
    let output = run_text_inference(&state, &prompt);

    let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(4);

    tokio::spawn(async move {
        let _ = tx
            .send(Ok(Event::default().data(
                serde_json::to_string(&StreamChunk {
                    delta: output,
                    done: true,
                })
                .unwrap(),
            )))
            .await;
    });

    Sse::new(ReceiverStream::new(rx))
}