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:
MechaCat02
2026-06-13 23:41:42 +02:00
parent 25aba3ac58
commit d3b827421f
6 changed files with 137 additions and 35 deletions

View File

@@ -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]