Update app.py

app.py

@@ -1,3 +1,12 @@
+To address the `ValueError` from the previous conversation, the primary issue is how the `_build_inputs` function handles the processor call, as it doesn't explicitly pass `max_length`. This can lead to the processor's internal logic truncating the input in a way that removes the image tokens, causing a mismatch.
+
+Here are the complete corrections for your `app.py` file to fix that issue and improve the overall code. The most important change is to the `_build_inputs` function.
+
+-----
+
+### Corrected Code for `app.py`
+
+````python
 # app.py — HTR Space with Feedback Loop, Memory Post-Correction, and GRPO Export
 
 import os, time, json, hashlib, difflib, uuid, csv
@@ -18,10 +27,10 @@ from jiwer import cer
 
 # ---------------- Storage & Paths ----------------
 os.makedirs("data", exist_ok=True)
-FEEDBACK_PATH = "data/feedback.jsonl"  # raw feedback log (per sample)
-MEMORY_RULES_PATH = "data/memory_rules.json"  # compiled post-correction rules
-GRPO_EXPORT_PATH = "data/grpo_prefs.jsonl"  # preference pairs for GRPO
-CSV_EXPORT_PATH = "data/feedback.csv"  # optional tabular export
+FEEDBACK_PATH = "data/feedback.jsonl" # raw feedback log (per sample)
+MEMORY_RULES_PATH = "data/memory_rules.json" # compiled post-correction rules
+GRPO_EXPORT_PATH = "data/grpo_prefs.jsonl" # preference pairs for GRPO
+CSV_EXPORT_PATH = "data/feedback.csv" # optional tabular export
 
 # ---------------- Models ----------------
 MODEL_PATHS = {
@@ -69,10 +78,27 @@ def warmup(progress=gr.Progress(track_tqdm=True)):
 # ---------------- Helpers ----------------
 def _build_inputs(processor, tokenizer, image: Image.Image, prompt: str):
     messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}]
+    
+    # We explicitly set max_length and truncation here to resolve the token mismatch error.
+    # A value of 2048 is safe, as an image takes up ~1024 tokens.
+    max_len_val = 2048
+
     if tokenizer and hasattr(tokenizer, "apply_chat_template"):
         chat_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        return processor(text=[chat_prompt], images=[image], return_tensors="pt")
-    return processor(text=[prompt], images=[image], return_tensors="pt")
+        return processor(
+            text=[chat_prompt],
+            images=[image],
+            return_tensors="pt",
+            max_length=max_len_val,
+            truncation=True
+        )
+    return processor(
+        text=[prompt],
+        images=[image],
+        return_tensors="pt",
+        max_length=max_len_val,
+        truncation=True
+    )
 
 def _decode_text(model, processor, tokenizer, output_ids, prompt: str):
     try:
@@ -161,7 +187,7 @@ def _compile_rules_from_feedback(min_count: int = 2, max_phrase_len: int = 40):
                 row = json.loads(line)
             except Exception:
                 continue
-            if row.get("reward", 0) < 1:  # only learn from thumbs-up or explicit 'accepted_correction'
+            if row.get("reward", 0) < 1: # only learn from thumbs-up or explicit 'accepted_correction'
                 continue
             pred = _safe_text(row.get("prediction", ""))
             corr = _safe_text(row.get("correction", "")) or _safe_text(row.get("ground_truth", ""))
@@ -209,7 +235,7 @@ def ocr_image(image: Image.Image, model_choice: str, query: str = None,
     batch = _build_inputs(processor, tokenizer, image, prompt).to(device)
     with torch.inference_mode():
         output_ids = model.generate(**batch, max_new_tokens=max_new_tokens, do_sample=False,
-                                    temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty)
+                                     temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty)
     raw = _decode_text(model, processor, tokenizer, output_ids, prompt).replace("<|im_end|>", "").strip()
     # Apply memory post-correction
     post = _apply_memory(raw, model_choice, use_memory)
@@ -391,7 +417,7 @@ from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from trl import GRPOConfig, GRPOTrainer
 
-MODEL_ID = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-VL-7B-Instruct")  # change if needed
+MODEL_ID = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-VL-7B-Instruct") # change if needed
 OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "grpo_output")
 DATA_PATH = os.environ.get("DATA_PATH", "data/grpo_prefs.jsonl")
 
@@ -443,7 +469,7 @@ def main():
 
     trainer = GRPOTrainer(
         model=model,
-        ref_model=None,  # let TRL create a frozen copy internally
+        ref_model=None, # let TRL create a frozen copy internally
         args=cfg,
         tokenizer=tok,
         train_dataset=ds
@@ -606,4 +632,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
 # The `if __name__ == "__main__":` block should be at the top level
 if __name__ == "__main__":
-    demo.queue(max_size=50).launch(share=True)
+    demo.queue(max_size=50).launch(share=True)
+````