Ensure all samplers are freed on error

Signed-off-by: Adrien Gallouët <angt@huggingface.co>
This commit is contained in:
Adrien Gallouët 2025-02-05 15:42:59 +00:00
parent 5b777877b1
commit 695b1292e9
No known key found for this signature in database

View File

@ -365,14 +365,17 @@ impl LlamacppSampler {
let dist = unsafe { let dist = unsafe {
llamacpp::sampler_init_dist(req.seed) llamacpp::sampler_init_dist(req.seed)
}; };
let all = &[
("top_k", top_k),
("top_p", top_p),
("typical_p", typical_p),
("temp", temp),
("penalties", penalties),
("dist", dist),
];
let mut failed = false; let mut failed = false;
for (k, v) in &[( "top_k", top_k ), for (k, v) in all {
( "top_p", top_p ),
("typical_p", typical_p),
( "temp", temp ),
("penalties", penalties),
( "dist", dist )] {
if v.is_null() { if v.is_null() {
error!("Failed to init {k} sampler"); error!("Failed to init {k} sampler");
failed = true; failed = true;
@ -381,6 +384,7 @@ impl LlamacppSampler {
} }
} }
if failed { if failed {
unsafe { llamacpp::sampler_free(chain) };
None None
} else { } else {
Some(LlamacppSampler{chain}) Some(LlamacppSampler{chain})