fix(analysis): guardrails against model repetition loops
Small vision models (qwen3-vl-4b) sometimes loop the OCR array, running the response into the token ceiling (finish_reason: length) and failing to parse. Layered guardrails: - Prompts: explicit "transcribe each element once, never repeat/loop, at most N, STOP when done" in the OCR, grounding and combined prompts. - Schema: hard maxItems on ocr_results/tagging_results/content_type and maxLength on text/scene — LM Studio's grammar enforces these, so a loop is grammar-bounded rather than relying on max_tokens. - Sampling: a configurable frequency_penalty (ANALYSIS_FREQUENCY_PENALTY, default 0.3) sent with each request — the decode-time lever that actually breaks loops. - Code: sanitize() collapses runaway consecutive duplicates (same normalized text + kind) so a loop that slips through still can't flood the row. Tests: schema caps present, frequency_penalty included only when nonzero, sanitize collapses consecutive repeats; config default. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -36,6 +36,7 @@ pub struct VisionClient {
|
||||
api_key: Option<String>,
|
||||
max_tokens: u32,
|
||||
response_format: ResponseFormat,
|
||||
frequency_penalty: f64,
|
||||
slice: SliceParams,
|
||||
}
|
||||
|
||||
@@ -58,6 +59,7 @@ impl VisionClient {
|
||||
api_key: cfg.api_key.clone(),
|
||||
max_tokens: cfg.max_tokens,
|
||||
response_format: cfg.response_format,
|
||||
frequency_penalty: cfg.frequency_penalty,
|
||||
slice: SliceParams {
|
||||
max_pixels: cfg.max_pixels,
|
||||
min_slice_height: cfg.min_slice_height,
|
||||
@@ -75,8 +77,13 @@ impl VisionClient {
|
||||
// and let the server cope (preserves the prior behavior).
|
||||
let Some(img) = image::load_from_memory(image).ok() else {
|
||||
let url = format!("data:{mime};base64,{}", b64(image));
|
||||
let body =
|
||||
build_request_body(&self.model, self.max_tokens, &url, self.response_format);
|
||||
let body = build_request_body(
|
||||
&self.model,
|
||||
self.max_tokens,
|
||||
&url,
|
||||
self.response_format,
|
||||
self.frequency_penalty,
|
||||
);
|
||||
return parse_chat_completion(&self.post_chat(body).await?);
|
||||
};
|
||||
|
||||
@@ -89,6 +96,7 @@ impl VisionClient {
|
||||
self.max_tokens,
|
||||
&data_url(&jpeg),
|
||||
self.response_format,
|
||||
self.frequency_penalty,
|
||||
);
|
||||
parse_chat_completion(&self.post_chat(body).await?)
|
||||
}
|
||||
@@ -115,6 +123,7 @@ impl VisionClient {
|
||||
&self.model,
|
||||
self.max_tokens,
|
||||
self.response_format,
|
||||
self.frequency_penalty,
|
||||
&data_url(&jpeg),
|
||||
);
|
||||
let parsed = parse_chat_completion(&self.post_chat(body).await?)?;
|
||||
@@ -138,6 +147,7 @@ impl VisionClient {
|
||||
&self.model,
|
||||
self.max_tokens,
|
||||
self.response_format,
|
||||
self.frequency_penalty,
|
||||
&data_url(&whole),
|
||||
&ocr_text,
|
||||
);
|
||||
@@ -446,11 +456,13 @@ pub fn build_request_body(
|
||||
max_tokens: u32,
|
||||
data_url: &str,
|
||||
response_format: ResponseFormat,
|
||||
frequency_penalty: f64,
|
||||
) -> serde_json::Value {
|
||||
build_body(
|
||||
model,
|
||||
max_tokens,
|
||||
response_format,
|
||||
frequency_penalty,
|
||||
prompt::SYSTEM_PROMPT,
|
||||
"page_analysis",
|
||||
prompt::output_json_schema(),
|
||||
@@ -464,12 +476,14 @@ fn build_ocr_body(
|
||||
model: &str,
|
||||
max_tokens: u32,
|
||||
response_format: ResponseFormat,
|
||||
frequency_penalty: f64,
|
||||
data_url: &str,
|
||||
) -> serde_json::Value {
|
||||
build_body(
|
||||
model,
|
||||
max_tokens,
|
||||
response_format,
|
||||
frequency_penalty,
|
||||
prompt::OCR_PROMPT,
|
||||
"page_ocr",
|
||||
prompt::ocr_json_schema(),
|
||||
@@ -483,6 +497,7 @@ fn build_grounding_body(
|
||||
model: &str,
|
||||
max_tokens: u32,
|
||||
response_format: ResponseFormat,
|
||||
frequency_penalty: f64,
|
||||
data_url: &str,
|
||||
ocr_text: &str,
|
||||
) -> serde_json::Value {
|
||||
@@ -491,6 +506,7 @@ fn build_grounding_body(
|
||||
model,
|
||||
max_tokens,
|
||||
response_format,
|
||||
frequency_penalty,
|
||||
prompt::GROUNDING_PROMPT,
|
||||
"page_grounding",
|
||||
prompt::grounding_json_schema(),
|
||||
@@ -500,11 +516,14 @@ fn build_grounding_body(
|
||||
}
|
||||
|
||||
/// Build a chat/completions body with a system prompt, an image part, an
|
||||
/// optional trailing text part, and the requested `response_format`.
|
||||
/// optional trailing text part, and the requested `response_format`. A
|
||||
/// non-zero `frequency_penalty` is included to suppress repetition loops.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn build_body(
|
||||
model: &str,
|
||||
max_tokens: u32,
|
||||
response_format: ResponseFormat,
|
||||
frequency_penalty: f64,
|
||||
system_prompt: &str,
|
||||
schema_name: &str,
|
||||
schema: serde_json::Value,
|
||||
@@ -525,6 +544,9 @@ fn build_body(
|
||||
{ "role": "user", "content": user_content }
|
||||
]
|
||||
});
|
||||
if frequency_penalty.abs() > f64::EPSILON {
|
||||
body["frequency_penalty"] = json!(frequency_penalty);
|
||||
}
|
||||
let rf = match response_format {
|
||||
ResponseFormat::None => None,
|
||||
ResponseFormat::JsonObject => Some(json!({ "type": "json_object" })),
|
||||
@@ -574,6 +596,22 @@ fn extract_json_object(s: &str) -> &str {
|
||||
/// content-warning values.
|
||||
pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
|
||||
a.ocr_results.retain(|r| !r.text.trim().is_empty());
|
||||
|
||||
// Collapse runaway repetition: a degenerating model emits the same line
|
||||
// over and over. Drop any piece identical (normalized text + kind) to the
|
||||
// one just kept, so a loop that slips past the schema caps still can't
|
||||
// flood the row.
|
||||
let mut prev: Option<String> = None;
|
||||
a.ocr_results.retain(|r| {
|
||||
let key = format!("{}\u{0}{}", r.kind, normalize_text(&r.text));
|
||||
if prev.as_deref() == Some(key.as_str()) {
|
||||
false
|
||||
} else {
|
||||
prev = Some(key);
|
||||
true
|
||||
}
|
||||
});
|
||||
|
||||
a.ocr_results.truncate(MAX_OCR_PIECES);
|
||||
for r in &mut a.ocr_results {
|
||||
r.text = truncate_chars(r.text.trim(), MAX_OCR_TEXT_CHARS);
|
||||
@@ -702,7 +740,8 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn combined_body_uses_full_schema_and_image() {
|
||||
let body = build_request_body("m", 100, "data:image/png;base64,AA", ResponseFormat::JsonSchema);
|
||||
let body =
|
||||
build_request_body("m", 100, "data:image/png;base64,AA", ResponseFormat::JsonSchema, 0.3);
|
||||
assert_eq!(body["messages"][0]["role"], "system");
|
||||
assert_eq!(
|
||||
body["messages"][1]["content"][0]["image_url"]["url"],
|
||||
@@ -714,18 +753,24 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ocr_body_carries_only_the_ocr_schema() {
|
||||
let body = build_ocr_body("m", 100, ResponseFormat::JsonSchema, "d");
|
||||
fn ocr_body_carries_only_the_ocr_schema_and_caps() {
|
||||
let body = build_ocr_body("m", 100, ResponseFormat::JsonSchema, 0.3, "d");
|
||||
let schema = &body["response_format"]["json_schema"]["schema"];
|
||||
assert_eq!(body["response_format"]["json_schema"]["name"], "page_ocr");
|
||||
let req = &body["response_format"]["json_schema"]["schema"]["required"];
|
||||
assert_eq!(req, &json!(["ocr_results"]));
|
||||
// No user text part — just the image.
|
||||
assert_eq!(schema["required"], json!(["ocr_results"]));
|
||||
// Hard repetition guards baked into the schema.
|
||||
assert!(schema["properties"]["ocr_results"]["maxItems"].is_number());
|
||||
assert!(
|
||||
schema["properties"]["ocr_results"]["items"]["properties"]["text"]["maxLength"]
|
||||
.is_number()
|
||||
);
|
||||
assert_eq!(body["messages"][1]["content"].as_array().unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grounding_body_includes_ocr_context_and_grounding_schema() {
|
||||
let body = build_grounding_body("m", 100, ResponseFormat::JsonSchema, "d", "[speech] Hi");
|
||||
let body =
|
||||
build_grounding_body("m", 100, ResponseFormat::JsonSchema, 0.3, "d", "[speech] Hi");
|
||||
assert_eq!(body["response_format"]["json_schema"]["name"], "page_grounding");
|
||||
let parts = body["messages"][1]["content"].as_array().unwrap();
|
||||
assert_eq!(parts.len(), 2, "image + text parts");
|
||||
@@ -736,10 +781,33 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn response_format_none_omits_the_field() {
|
||||
let body = build_request_body("m", 100, "d", ResponseFormat::None);
|
||||
let body = build_request_body("m", 100, "d", ResponseFormat::None, 0.3);
|
||||
assert!(body.get("response_format").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frequency_penalty_included_only_when_nonzero() {
|
||||
let on = build_ocr_body("m", 100, ResponseFormat::None, 0.4, "d");
|
||||
assert_eq!(on["frequency_penalty"], 0.4);
|
||||
let off = build_ocr_body("m", 100, ResponseFormat::None, 0.0, "d");
|
||||
assert!(off.get("frequency_penalty").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_collapses_runaway_repetition() {
|
||||
let body = chat_body(
|
||||
r#"{"ocr_results":[
|
||||
{"text":"I don't know","kind":"speech"},
|
||||
{"text":"I don't know","kind":"speech"},
|
||||
{"text":"I don't know","kind":"speech"},
|
||||
{"text":"ok","kind":"speech"}
|
||||
]}"#,
|
||||
);
|
||||
let v = parse_chat_completion(&body).unwrap();
|
||||
let texts: Vec<&str> = v.ocr_results.iter().map(|o| o.text.as_str()).collect();
|
||||
assert_eq!(texts, vec!["I don't know", "ok"], "consecutive repeats collapsed");
|
||||
}
|
||||
|
||||
// --- geometry: plan_slices ---
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user