sipp-rs 0.1.0

Unified Rust library for extensible Sipp inference
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

use std::collections::HashMap;

use crate::native_bridge::NativeRuntimeHandle;
use crate::runtime::config::NativeRuntimeConfig;
use crate::runtime::llama_seq_id;
use crate::runtime::scheduler::{SamplerCacheKey, SamplerHandle, SlotState};

use super::super::sampler::{attach_backend_sampler, create_sampler, ResidentBackendSampler};

pub(super) fn ensure_slot_sampler(
    slot: &mut SlotState,
    native_runtime: &mut NativeRuntimeHandle,
    config: &NativeRuntimeConfig,
    sampler_pool: &mut HashMap<SamplerCacheKey, Vec<SamplerHandle>>,
    resident_backend_samplers: &mut HashMap<llama_seq_id, ResidentBackendSampler>,
) -> bool {
    let Some(request) = slot.request() else {
        return false;
    };

    let sampling_json = match config.try_sampling_json_with_override(request.sampling.as_ref()) {
        Ok(sampling_json) => sampling_json,
        Err(error) => {
            slot.fail(format!(
                "Failed to serialize sampler configuration: {error}"
            ));
            return false;
        }
    };
    let key = SamplerCacheKey {
        sampling_json,
        grammar: request.grammar.clone(),
        json_schema: request.json_schema.clone(),
    };

    if let Some(resident) = resident_backend_samplers.remove(&slot.seq_id) {
        if resident.key == key {
            slot.set_sampler(resident.sampler);
            slot.sampler_key = Some(key);
            slot.backend_sampler_attached = true;
            return true;
        }

        if slot.seq_id >= 0 {
            native_runtime.detach_sampler(slot.seq_id);
        }
    }

    if let Some(sampler) = sampler_pool.get_mut(&key).and_then(|vec| vec.pop()) {
        slot.set_sampler(sampler);
        slot.sampler_key = Some(key);
        attach_backend_sampler(native_runtime, slot);
        return true;
    }

    let sampling = slot.request().and_then(|request| request.sampling.as_ref());
    match create_sampler(
        native_runtime,
        config,
        sampling,
        Some(&key.grammar),
        Some(&key.json_schema),
    ) {
        Ok(sampler) => {
            slot.set_sampler(sampler);
            slot.sampler_key = Some(key);
            attach_backend_sampler(native_runtime, slot);
            true
        }
        Err(_) => {
            let message = if key.grammar.is_empty() {
                "Failed to create per-slot sampler."
            } else {
                "Failed to create per-slot grammar sampler."
            };
            slot.fail(message);
            false
        }
    }
}