Added live text, show extraction, add website

app.py

@@ -1,5 +1,9 @@
 # Copyright: Shayekh Bin Islam. KAIST, South Korea. 2026.
 
+MAX_TEXT_CHAR = 1500
+
+# model_id = "Qwen/Qwen3.5-9B"
+model_id = "Qwen/Qwen3.5-2B"
 
 try:
     import spaces
@@ -54,6 +58,53 @@ def extract_pdf_content(pdf_path, max_pages=2):
         images.append(img)
     return text, images
 
+def extract_website_content(url, max_images=2):
+    """Extract text and images from a website URL."""
+    import requests
+    from bs4 import BeautifulSoup
+    import io
+    
+    headers = {
+        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
+    }
+    response = requests.get(url, headers=headers, timeout=10)
+    response.raise_for_status()
+    
+    soup = BeautifulSoup(response.content, 'html.parser')
+    
+    for script in soup(["script", "style", "nav", "footer", "header", "noscript"]):
+        script.extract()
+        
+    text = soup.get_text(separator='\n')
+    lines = (line.strip() for line in text.splitlines())
+    chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
+    text = '\n'.join(chunk for chunk in chunks if chunk)
+    
+    images = []
+    img_tags = soup.find_all('img')
+    for img in img_tags:
+        if len(images) >= max_images:
+            break
+        src = img.get('src') or img.get('data-src')
+        if src:
+            if src.startswith('//'):
+                src = 'https:' + src
+            elif src.startswith('/'):
+                from urllib.parse import urljoin
+                src = urljoin(url, src)
+                
+            try:
+                img_resp = requests.get(src, headers=headers, timeout=5)
+                if img_resp.status_code == 200:
+                    pil_img = Image.open(io.BytesIO(img_resp.content))
+                    if pil_img.mode != 'RGB':
+                        pil_img = pil_img.convert('RGB')
+                    if pil_img.width >= 100 and pil_img.height >= 100:
+                        images.append(pil_img)
+            except Exception as e:
+                print(f"Failed to load image {src}: {e}")
+                
+    return text, images
 
 def get_base64_image(image):
     buffered = io.BytesIO()
@@ -62,13 +113,13 @@ def get_base64_image(image):
     return f"data:image/jpeg;base64,{img_str}"
 
 @spaces.GPU(duration=120)
-def extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_lang):
+def extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_lang, max_text_char=1500):
     """Use Transformers to extract vocabulary from text and images."""
     global model, processor
     
     os.makedirs("log", exist_ok=True)
     
-    if pdf_text.strip() == "":
+    if len(pdf_text.strip()) == 0:
         pdf_text = '''"No Text available, see provided images only."'''
 
 
@@ -76,16 +127,20 @@ def extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_
     if translit_lang.upper() != "ENGLISH":
         non_english = f" CRITICAL: You MUST use the native alphabet/script of {translit_lang.upper()}, do NOT use English letters unless requested."
     
-    prompt_text = f"""Extract 3 to 5 key Korean words or phrases from the following text and images.
+    prompt_text = f"""Extract at least 10 key Korean words or phrases from the following text and images.
+Focus on meaningful vocabulary that is highly helpful for a new language learner (e.g., common nouns, verbs, adjectives, or useful expressions).
+CRITICAL: Do NOT extract website template words, navigation menus, boilerplate text, UI elements, or titles like 'Home page', 'News', 'Menu'.
+
 Return ONLY a valid JSON list of dictionaries, where each dictionary has four keys:
 - 'korean' (the Korean text)
 - 'transliteration' (the pronunciation transliterated into {translit_lang.upper()} script/characters, formatted as {translit_format}.{non_english})
 - 'translation' (the translation into {target_lang.upper()})
 - 'explanation' (a brief grammar or context note in {target_lang.upper()}).
 No markdown formatting, just raw JSON with ```json and ``` markers.
+CRITICAL: Do NOT provide any conversational filler, thinking steps, or reasoning. Answer quick without very long thinking. Output the JSON array IMMEDIATELY.
 
 Text:
-{pdf_text[:1500]}
+{pdf_text[:int(max_text_char)]}
 """
     
     # DEBUG: Log prompt text
@@ -122,13 +177,112 @@ Text:
             return_tensors="pt",
             padding=True
         ).to("cuda")
+        
+        from transformers import TextIteratorStreamer
+        from threading import Thread
 
-        generated_ids = model.generate(
+        streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(
             **inputs,
-            # max_new_tokens=2048*4,
+            streamer=streamer,
             max_new_tokens=2048*16,
-            temperature=1.0,
-            top_p=0.95,
+            do_sample=True,
+            repetition_penalty=1.1,
+        )
+
+        if len(images) > 0:
+            generation_kwargs.update(dict(temperature=0.6, top_p=0.95, top_k=20, min_p=0.0))
+        else:
+            generation_kwargs.update(dict(temperature=1.0, top_p=0.95, top_k=20, min_p=0.0))
+
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        output_text = ""
+        for new_text in streamer:
+            output_text += new_text
+            yield output_text, None
+
+        # DEBUG: Log raw output text
+        with open("log/debug_vlm_output.txt", "w", encoding="utf-8") as f:
+            f.write(output_text)
+            
+    except Exception as e:
+        print(f"Error during Transformers inference: {e}")
+        yield f"Error during Transformers inference: {e}", []
+        return
+
+    try:
+        import re
+        # Extract JSON from markdown code fences or raw output
+        json_matches = list(re.finditer(r'```(?:json)?\s*([\s\S]*?)```', output_text))
+        if json_matches:
+            clean_text = json_matches[-1].group(1).strip()
+        else:
+            # Fallback: find last [ ... ] or { ... } block
+            json_matches = list(re.finditer(r'(\[[\s\S]*\]|\{[\s\S]*\})', output_text))
+            clean_text = json_matches[-1].group(1).strip() if json_matches else output_text.strip()
+            
+        data = json.loads(clean_text)
+        if not isinstance(data, list):
+            data = [data]
+        yield output_text, data
+    except Exception as e:
+        print(f"Error parsing JSON: {e}\nRaw output: {output_text}")
+        yield output_text, []
+
+def translate_vocabulary(korean_words, translit_lang, translit_format, target_lang):
+    """Use Transformers text-only inference to translate/transliterate Korean words."""
+    global model, processor
+    
+    non_english = ""
+    if translit_lang.upper() != "ENGLISH":
+        non_english = f" CRITICAL: You MUST use the native alphabet/script of {translit_lang.upper()}, do NOT use English letters unless requested."
+    
+    words_str = ", ".join(korean_words)
+    prompt_text = f"""Translate and transliterate the following Korean words.
+Return ONLY a valid JSON list of dictionaries, where each dictionary has four keys:
+- 'korean' (the original Korean text)
+- 'transliteration' (the pronunciation transliterated into {translit_lang.upper()} script/characters, formatted as {translit_format}.{non_english})
+- 'translation' (the translation into {target_lang.upper()})
+- 'explanation' (a brief grammar or context note in {target_lang.upper()}).
+No markdown formatting, just raw JSON with ```json and ``` markers.
+CRITICAL: Do NOT provide any conversational filler, thinking steps, or reasoning. Answer quick without very long thinking. Output the JSON array IMMEDIATELY.
+
+Korean words:
+{words_str}
+"""
+
+    # DEBUG: Log translation prompt text
+    with open("log/debug_translate_prompt.txt", "w", encoding="utf-8") as f:
+        f.write(prompt_text)
+
+    messages = [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": prompt_text}]
+        }
+    ]
+
+    try:
+        model.to("cuda")
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(
+            text=[text],
+            images=None,
+            return_tensors="pt",
+            padding=True
+        ).to("cuda")
+
+        generated_ids = model.generate(
+            **inputs,
+            # max_new_tokens=2048*16,
+            max_new_tokens=2048*2,
+            # temperature=1.0,
+            # top_p=0.95,
+            temperature=1.0, top_p=0.95, top_k=20, min_p=0.0, 
+            # presence_penalty=1.5, 
+            repetition_penalty=1.1,
             do_sample=True
         )
         
@@ -137,24 +291,22 @@ Text:
         ]
         output_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         
-        # DEBUG: Log raw output text
-        with open("log/debug_vlm_output.txt", "w", encoding="utf-8") as f:
+        # DEBUG: Log raw translation output text
+        with open("log/debug_translate_output.txt", "w", encoding="utf-8") as f:
             f.write(output_text)
             
     except Exception as e:
-        print(f"Error during Transformers inference: {e}")
+        print(f"Error during Transformers text inference: {e}")
         return []
 
     try:
         import re
-        # Extract JSON from markdown code fences or raw output
-        json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', output_text)
-        if json_match:
-            clean_text = json_match.group(1).strip()
+        json_matches = list(re.finditer(r'```(?:json)?\s*([\s\S]*?)```', output_text))
+        if json_matches:
+            clean_text = json_matches[-1].group(1).strip()
         else:
-            # Fallback: find first [ ... ] or { ... } block
-            json_match = re.search(r'(\[[\s\S]*\]|\{[\s\S]*\})', output_text)
-            clean_text = json_match.group(1).strip() if json_match else output_text.strip()
+            json_matches = list(re.finditer(r'(\[[\s\S]*\]|\{[\s\S]*\})', output_text))
+            clean_text = json_matches[-1].group(1).strip() if json_matches else output_text.strip()
             
         data = json.loads(clean_text)
         if not isinstance(data, list):
@@ -172,8 +324,14 @@ def numpy_to_base64_audio(wav, sample_rate):
     audio_base64 = base64.b64encode(buffer.read()).decode('utf-8')
     return f"data:audio/wav;base64,{audio_base64}"
 
+import hashlib
+
+def hash_file(filepath):
+    with open(filepath, 'rb') as f:
+        return hashlib.md5(f.read(1024*1024)).hexdigest()
+
 @spaces.GPU(duration=120)
-def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=gr.Progress()):
+def process_pdf(pdf_file, url_input, translit_lang, translit_format, target_lang, max_text_char, last_source_hash, last_korean_words, progress=gr.Progress()):
     global tts, voice_style
     
     # Clean language choices from "Family - Language" to just "Language"
@@ -184,26 +342,59 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
         
     os.makedirs("log", exist_ok=True)
     
-    if pdf_file is None:
-        return "<p>Please upload a PDF.</p>"
+    is_url = bool(url_input and url_input.strip())
+    if pdf_file is None and not is_url:
+        yield "<p>Please upload a PDF or enter a URL.</p>", None, None, ""
+        return
+        
+    if is_url:
+        current_source_hash = hashlib.md5(url_input.strip().encode()).hexdigest()
+    else:
+        current_source_hash = hash_file(pdf_file.name)
+        
+    vocab_list = []
+
+    # if last_source_hash == current_source_hash and last_korean_words:
+    #     # Just run text-to-text LLM
+    #     progress(0.2, desc="Translating previously extracted vocabulary...")
+    #     korean_words = [item.get("korean") for item in last_korean_words if item.get("korean")]
+    #     for attempt in range(1, 4):
+    #         vocab_list = translate_vocabulary(korean_words, translit_lang, translit_format, target_lang)
+    #         if vocab_list:
+    #             break
+    # else:
     
     try:
-        progress(0, desc="Reading PDF...")
-        pdf_text, images = extract_pdf_content(pdf_file.name)
-        if not pdf_text.strip() and not images:
-            return "<p>No content found in PDF.</p>"
+        if is_url:
+            progress(0, desc="Fetching Website...")
+            content_text, images = extract_website_content(url_input.strip())
+        else:
+            progress(0, desc="Reading PDF...")
+            content_text, images = extract_pdf_content(pdf_file.name)
+            
+        if not content_text.strip() and not images:
+            yield "<p>No content found.</p>", current_source_hash, None, "", "", []
+            return
     except Exception as e:
-        return f"<p>Error reading PDF: {e}</p>"
+        yield f"<p>Error reading content: {e}</p>", None, None, "", "", []
+        return
 
     vocab_list = []
+    stream_text = ""
     for attempt in range(1, 4):
         progress(0.2, desc=f"Extracting vocabulary (Attempt {attempt}/3)...")
-        vocab_list = extract_vocabulary(pdf_text, images, translit_lang, translit_format, target_lang)
+        for stream_t, v_list in extract_vocabulary(content_text, images, translit_lang, translit_format, target_lang, max_text_char):
+            stream_text = stream_t
+            if v_list is not None:
+                vocab_list = v_list
+            yield "", current_source_hash, None, stream_text, content_text, images
+            
         if vocab_list:
             break
 
     if not vocab_list:
-        return "<p>Failed to extract vocabulary after 3 attempts. The model might not have found Korean text or returned an invalid format.</p>"
+        yield "<p>Failed to extract or translate vocabulary after 3 attempts.</p>", current_source_hash, None, stream_text, content_text, images
+        return
 
     progress(0.6, desc="Generating TTS audio...")
     # Pre-generate TTS audio
@@ -214,7 +405,11 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
             korean += "."
             
         try:
-            wav, dur = tts.synthesize(korean, voice_style=voice_style, lang="ko")
+            wav, dur = tts.synthesize(
+                korean, voice_style=voice_style, lang="ko",
+                total_steps=12,
+                speed=0.7,
+            )
             
             # DEBUG: Save audio locally
             wav_1d = wav.squeeze()
@@ -250,8 +445,8 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
         }}
         .flashcard {{
             width: 100%;
-            height: 350px;
-            position: relative;
+            min-height: 400px;
+            display: grid;
             transition: transform 0.6s cubic-bezier(0.4, 0.2, 0.2, 1);
             transform-style: preserve-3d;
             cursor: pointer;
@@ -260,9 +455,8 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
             transform: rotateY(180deg);
         }}
         .card-face {{
-            position: absolute;
+            grid-area: 1 / 1;
             width: 100%;
-            height: 100%;
             backface-visibility: hidden;
             display: flex;
             flex-direction: column;
@@ -370,6 +564,16 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
             font-size: 14px;
             font-weight: 600;
         }}
+        .copy-icon {{
+            cursor: pointer;
+            color: #a0a0a0;
+            font-size: 24px;
+            transition: color 0.2s, transform 0.2s;
+        }}
+        .copy-icon:hover {{
+            color: #7c3aed;
+            transform: scale(1.1);
+        }}
     </style>
     </head>
     <body>
@@ -377,12 +581,18 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
         <div class="flashcard-container">
             <div class="flashcard" id="card" onclick="flipCard()">
                 <div class="card-face card-front">
-                    <div class="korean-text" id="front-text"><i class="fi fi-rr-spinner-third fa-spin"></i> Loading...</div>
+                    <div style="display: flex; align-items: center; justify-content: center; gap: 15px; margin-bottom: 20px; width: 100%;">
+                        <div class="korean-text" id="front-text" style="margin-bottom: 0;"><i class="fi fi-rr-spinner-third fa-spin"></i> Loading...</div>
+                        <i class="fi fi-rr-copy copy-icon" id="copy-ko" onclick="copyText('front-text', event, this)" title="Copy text" style="display:none;"></i>
+                    </div>
                     <button class="audio-btn" onclick="playAudio(event)" id="audio-btn" style="display:none;"><i class="fi fi-rr-play-circle"></i> Play Audio</button>
                     <p style="margin-top:20px; color:#999; font-size:13px; display:flex; align-items:center; gap:5px;"><i class="fi fi-rr-rotate-right"></i> Click card to flip 🎯</p>
                 </div>
                 <div class="card-face card-back">
-                    <div class="english-text" id="back-en"></div>
+                    <div style="display: flex; align-items: center; justify-content: center; gap: 15px; margin-bottom: 5px; width: 100%;">
+                        <div class="english-text" id="back-en" style="margin-bottom: 0;"></div>
+                        <i class="fi fi-rr-copy copy-icon" id="copy-en" onclick="copyText('back-en', event, this)" title="Copy text" style="display:none;"></i>
+                    </div>
                     <div class="translit-text" id="back-translit"></div>
                     <div class="explanation-text"><i class="fi fi-rr-lightbulb-on" style="color:#f1c40f;"></i> <span id="back-exp"></span></div>
                 </div>
@@ -405,6 +615,8 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
                 document.getElementById('front-text').innerHTML = "No vocabulary found 😥";
                 document.getElementById('prev-btn').disabled = true;
                 document.getElementById('next-btn').disabled = true;
+                if(document.getElementById('copy-ko')) document.getElementById('copy-ko').style.display = 'none';
+                if(document.getElementById('copy-en')) document.getElementById('copy-en').style.display = 'none';
                 return;
             }}
             const card = cards[currentIndex];
@@ -413,6 +625,9 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
             document.getElementById('back-translit').innerText = card.transliteration ? `[${{card.transliteration}}]` : "";
             document.getElementById('back-exp').innerText = card.explanation || "";
             
+            if(document.getElementById('copy-ko')) document.getElementById('copy-ko').style.display = 'block';
+            if(document.getElementById('copy-en')) document.getElementById('copy-en').style.display = 'block';
+            
             document.getElementById('prev-btn').disabled = currentIndex === 0;
             document.getElementById('next-btn').disabled = currentIndex === cards.length - 1;
             document.getElementById('progress-text').innerHTML = `📚 Card ${{currentIndex + 1}} of ${{cards.length}}`;
@@ -452,6 +667,20 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
             }}
         }}
 
+        function copyText(elementId, e, iconEl) {{
+            e.stopPropagation();
+            const textToCopy = document.getElementById(elementId).innerText;
+            navigator.clipboard.writeText(textToCopy).then(() => {{
+                const oldClass = iconEl.className;
+                iconEl.className = "fi fi-rr-check copy-icon";
+                iconEl.style.color = "#27ae60";
+                setTimeout(() => {{
+                    iconEl.className = oldClass;
+                    iconEl.style.color = "";
+                }}, 1500);
+            }}).catch(err => console.log('Copy failed', err));
+        }}
+
         window.onload = function() {{
             updateCard();
         }};
@@ -464,7 +693,7 @@ def process_pdf(pdf_file, translit_lang, translit_format, target_lang, progress=
     safe_srcdoc = html.escape(iframe_html)
     
     # Return the iframe containing the whole SPA
-    return f'<iframe srcdoc="{safe_srcdoc}" style="width: 100%; height: 500px; border: none; overflow: hidden;"></iframe>'
+    yield f'<iframe srcdoc="{safe_srcdoc}" style="width: 100%; height: 650px; border: none; overflow-y: auto;"></iframe>', current_source_hash, vocab_list, stream_text, content_text, images
 
 LANGUAGE_DATA = """Indo-European	English, French, Portuguese, German, Romanian, Swedish, Danish, Bulgarian, Russian, Czech, Greek, Ukrainian, Spanish, Dutch, Slovak, Croatian, Polish, Lithuanian, Norwegian Bokmål, Norwegian Nynorsk, Persian, Slovenian, Gujarati, Latvian, Italian, Occitan, Nepali, Marathi, Belarusian, Serbian, Luxembourgish, Venetian, Assamese, Welsh, Silesian, Asturian, Chhattisgarhi, Awadhi, Maithili, Bhojpuri, Sindhi, Irish, Faroese, Hindi, Punjabi, Bengali, Oriya, Tajik, Eastern Yiddish, Lombard, Ligurian, Sicilian, Friulian, Sardinian, Galician, Catalan, Icelandic, Tosk Albanian, Limburgish, Dari, Afrikaans, Macedonian, Sinhala, Urdu, Magahi, Bosnian, Armenian, Latgalian, Scottish Gaelic, Central Kurdish, Northern Kurdish, Southern Pashto, Sanskrit, Dhundari, Marwari, Ahirani, Bagheli, Bagri, Bundeli, Braj, Kumaoni, Kashmiri
 Sino-Tibetan	Chinese (Simplified), Chinese (Traditional), Cantonese, Burmese, Standard Tibetan, Meitei
@@ -499,13 +728,22 @@ def get_example_pdf():
 def create_demo():
     example_pdf = get_example_pdf()
     
-    with gr.Blocks(title="LocalDuo") as demo:
-        gr.Markdown("# 🇰🇷✨ LocalDuo - Learn Korean from PDFs")
-        gr.Markdown("Upload a Korean book 📖 or document PDF 📄. The app uses **vLLM** 🧠 with Qwen3.5-2B to extract vocabulary from text and images, and Supertonic 🗣️ to generate pronunciation audio.")
+    css = """
+    #stream_box .progress-text, 
+    #stream_box .progress-level, 
+    #stream_box .progress, 
+    #stream_box .progress-container { 
+        display: none !important; 
+    }
+    """
+    with gr.Blocks(title="LocalDuo", css=css) as demo:
+        gr.Markdown("# 🇰🇷✨ LocalDuo - Learn Korean from PDFs & Websites")
+        gr.Markdown("Enter a website URL 🌐 or upload a Korean book PDF 📄. The app uses a **Vision-Language Model (VLM)** 🧠 to extract vocabulary from text and images, and a **Text-to-Speech (TTS)** engine 🗣️ to generate pronunciation audio.")
         
         with gr.Row():
             with gr.Column(scale=1):
-                pdf_input = gr.File(label="Upload Book PDF 📚", file_types=[".pdf"], value=example_pdf)
+                url_input = gr.Textbox(label="Enter a Website URL 🌐", placeholder="e.g. https://www.bbc.com/korean", value="https://www.bbc.com/korean")
+                pdf_input = gr.File(label="Or Upload Book PDF 📚", file_types=[".pdf"], value=example_pdf)
                 
                 gr.Markdown("### ⚙️ Customization Settings")
                 translit_lang = gr.Dropdown(
@@ -519,16 +757,42 @@ def create_demo():
                     choices=LANGUAGE_CHOICES, 
                     value="Indo-European - English"
                 )
+                max_text_char_input = gr.Slider(minimum=1000, maximum=30000, step=1000, value=1500, label="Max Input Text Length (Characters)")
                 
-                submit_btn = gr.Button("✨ Generate Flashcards ✨", variant="primary")
+                with gr.Row():
+                    submit_btn = gr.Button("✨ Generate Flashcards ✨", variant="primary")
+                    stop_btn = gr.Button("🛑 Stop Generation", variant="stop")
             
             with gr.Column(scale=2):
                 output_html = gr.HTML(label="Flashcards will appear here")
+                stream_box = gr.Textbox(label="Live Model Generation 🧠", lines=10, max_lines=20, interactive=False, autoscroll=True, elem_id="stream_box")
+                
+                with gr.Accordion("📄 Extracted Source Content", open=True):
+                    extracted_text_box = gr.Textbox(label="Extracted Text", lines=10, max_lines=15, interactive=False)
+                    extracted_images_gallery = gr.Gallery(label="Extracted Images", columns=4, height="auto", object_fit="contain")
                 
-        submit_btn.click(
+        last_source_state = gr.State(None)
+        last_korean_words_state = gr.State(None)
+
+        generate_event = submit_btn.click(
             fn=process_pdf, 
-            inputs=[pdf_input, translit_lang, translit_format, target_lang], 
-            outputs=output_html
+            inputs=[pdf_input, url_input, translit_lang, translit_format, target_lang, max_text_char_input, last_source_state, last_korean_words_state], 
+            outputs=[output_html, last_source_state, last_korean_words_state, stream_box, extracted_text_box, extracted_images_gallery]
+        )
+        
+        stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[generate_event])
+        
+        # Force autoscroll using Custom JS
+        stream_box.change(
+            fn=None,
+            js="""
+            function() {
+                const ta = document.querySelector('#stream_box textarea');
+                if (ta) {
+                    ta.scrollTop = ta.scrollHeight;
+                }
+            }
+            """
         )
     return demo
 
@@ -549,6 +813,162 @@ if __name__ == "__main__":
     #         processor.chat_template = f.read()
     # except Exception as e:
     #     print("Could not load custom chat template:", e)
+    processor.chat_template = """
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}
+""".strip()
         
     model = AutoModelForImageTextToText.from_pretrained(
         model_id,