Update index.html

changing text model

index.html

@@ -33,6 +33,11 @@
     <h2>Refined caption</h2>
     <p id="caption"></p>
 
+    <details>
+        <summary>Debug</summary>
+        <pre id="debug" style="white-space: pre-wrap;"></pre>
+    </details>
+
     <script type="module">
         import {
             Florence2ForConditionalGeneration,
@@ -48,9 +53,11 @@
         const status = document.getElementById("status");
         const rawCaptionEl = document.getElementById("rawCaption");
         const caption = document.getElementById("caption");
+        const debugEl = document.getElementById("debug");
 
         const MODEL_ID = "onnx-community/Florence-2-base-ft";
-        const REFINER_ID = "onnx-community/Qwen2.5-0.5B-Instruct";
+        // Text2text model: much more robust for rewrite tasks than a tiny chat LLM.
+        const REFINER_ID = "Xenova/LaMini-Flan-T5-248M";
 
         let selectedFile = null;
         let model = null;
@@ -58,15 +65,22 @@
         let tokenizer = null;
         let refiner = null;
 
+        function log(...args) {
+            console.log(...args);
+            debugEl.textContent += args.map(a =>
+                typeof a === "string" ? a : JSON.stringify(a)
+            ).join(" ") + "\n";
+        }
+
         // ---------- filename helpers ----------
         function cleanFilename(name) {
-            let base = name.replace(/\.[^.]+$/, "");                       // drop extension
-            base = base.replace(/[_-]+/g, " ");                            // _ and - -> space
+            let base = name.replace(/\.[^.]+$/, "");
+            base = base.replace(/[_-]+/g, " ");
             base = base.replace(
                 /\b(img|dsc|photo|image|picture|pic|file|screenshot|final|copy|edit|v\d+)\b/gi,
                 " "
-            );                                                              // drop generic tokens
-            base = base.replace(/\b\d+\b/g, " ");                          // drop bare numbers
+            );
+            base = base.replace(/\b\d+\b/g, " ");
             base = base.replace(/\s+/g, " ").trim();
             return base;
         }
@@ -94,45 +108,59 @@
         async function loadRefiner() {
             if (refiner) return;
             status.textContent = "Loading caption refiner...";
-            refiner = await pipeline("text-generation", REFINER_ID, {
-                dtype: "q4",
-                device: "webgpu",
+            refiner = await pipeline("text2text-generation", REFINER_ID, {
+                dtype: "q8",
+                device: "wasm", // safer than webgpu for this model
             });
         }
 
+        // ---------- sanity checks ----------
+        function looksBroken(text, rawCaption) {
+            if (!text) return true;
+            const t = text.trim();
+            if (t.length < 5) return true;
+            if (t.length > 220) return true;                       // runaway generation
+            if (/\b(\w+)\s+\1\s+\1\b/i.test(t)) return true;       // "the the the"
+            // If it contains words like "Raw caption:" or "Filename hints:", it echoed the prompt
+            if (/raw caption|filename hints|rewritten caption/i.test(t)) return true;
+            // If it is basically the raw caption, prefer the raw caption (no loss)
+            const norm = s => s.toLowerCase().replace(/[^a-z0-9 ]/g, "").trim();
+            if (norm(t) === norm(rawCaption)) return true;
+            return false;
+        }
+
         // ---------- refinement ----------
         async function refine(rawCaption, filename) {
             const hints = cleanFilename(filename);
 
-            const system =
-                "You rewrite image captions into ONE natural, fluent sentence. " +
-                "Use the filename hints only if they add useful information for describing " +
-                "the image; ignore them if they are generic or irrelevant. " +
-                "Do not invent details that are not in the caption or hints. " +
-                "Output only the rewritten sentence, nothing else.";
-
-            const user = hints
-                ? `Raw caption: ${rawCaption}\nFilename hints: ${hints}\nRewritten caption:`
-                : `Raw caption: ${rawCaption}\n(No useful filename hints.)\nRewritten caption:`;
-
-            const out = await refiner(
-                [
-                    { role: "system", content: system },
-                    { role: "user", content: user },
-                ],
-                { max_new_tokens: 60, do_sample: false }
-            );
+            // Flan-T5 wants a plain string prompt, and explicit instructions
+            // about the output format work best.
+            const prompt = hints
+                ? `Rewrite the caption into one fluent natural sentence. ` +
+                  `Use the filename hints only if they help describe the image; ` +
+                  `otherwise ignore them. Do not add details that are not present. ` +
+                  `Caption: ${rawCaption} ` +
+                  `Filename hints: ${hints} ` +
+                  `Rewritten caption:`
+                : `Rewrite the caption into one fluent natural sentence. ` +
+                  `Do not add details that are not present. ` +
+                  `Caption: ${rawCaption} ` +
+                  `Rewritten caption:`;
+
+            log("[refine prompt]", prompt);
+
+            const out = await refiner(prompt, {
+                max_new_tokens: 60,
+                do_sample: false,
+                num_beams: 4,
+                repetition_penalty: 1.3,
+                no_repeat_ngram_size: 3,
+            });
 
-            const generated = out[0].generated_text;
+            log("[refine raw output]", JSON.stringify(out));
 
-            // transformers.js returns either a string or the full messages array
-            let text;
-            if (Array.isArray(generated)) {
-                text = generated[generated.length - 1].content;
-            } else {
-                text = generated;
-            }
-            return (text || "").trim();
+            const text = (out?.[0]?.generated_text || "").trim();
+            return text;
         }
 
         // ---------- file input ----------
@@ -149,6 +177,7 @@
             caption.textContent = "";
             rawCaptionEl.textContent = "";
             status.textContent = "";
+            debugEl.textContent = "";
         });
 
         // ---------- main ----------
@@ -157,6 +186,7 @@
                 captionButton.disabled = true;
                 caption.textContent = "";
                 rawCaptionEl.textContent = "";
+                debugEl.textContent = "";
 
                 await loadCaptionModel();
                 status.textContent = "Preparing image...";
@@ -183,19 +213,30 @@
                 );
                 const rawCaption = result[task] || "";
 
+                log("[raw caption]", rawCaption);
+                log("[filename]", selectedFile.name);
+                log("[cleaned filename]", cleanFilename(selectedFile.name));
+
                 rawCaptionEl.textContent = rawCaption || "(no caption)";
 
                 await loadRefiner();
                 status.textContent = "Refining caption...";
                 const refined = await refine(rawCaption, selectedFile.name);
 
-                caption.textContent =
-                    refined || rawCaption || "No caption generated.";
-                status.textContent = "Done.";
+                log("[refined]", refined);
+
+                if (looksBroken(refined, rawCaption)) {
+                    log("[fallback] refined output rejected, using raw caption");
+                    caption.textContent = rawCaption || "No caption generated.";
+                    status.textContent = "Done (kept raw caption).";
+                } else {
+                    caption.textContent = refined;
+                    status.textContent = "Done.";
+                }
             } catch (error) {
                 console.error(error);
-                status.textContent =
-                    "Error: " + (error.message || error);
+                status.textContent = "Error: " + (error.message || error);
+                log("[error]", String(error));
             } finally {
                 captionButton.disabled = false;
             }