sync QWENVL-PRIVATE from /opt/tiger/thothvl_pretrain (HF tokens redacted)

.gitattributes

@@ -33,3 +33,27 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/agent_function_call/mobile_en_example.png filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/agent_function_call/mobile_zh_example.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/computer_use/computer_use1.jpeg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/computer_use/computer_use2.jpeg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example1.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example2.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example3.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example4.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example5.png filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example6.png filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example7.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/document_parsing/docparsing_example8.png filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/ocr/ocr_example1.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/ocr/ocr_example3.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/ocr/ocr_example5.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/spatial_understanding/Origamis.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/spatial_understanding/cakes.png filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/universal_recognition/unireco_bird_example.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/universal_recognition/unireco_birds_example.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/universal_recognition/unireco_celebrities_example.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/cookbooks/assets/universal_recognition/unireco_landmarks_example.jpg filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/qwen-vl-finetune/demo/videos/v_7bUu05RIksU.mp4 filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/qwen-vl-finetune/demo/videos/v_TpB_zMG3XBA.mp4 filter=lfs diff=lfs merge=lfs -text
+QWENVL-PRIVATE/qwen-vl-finetune/demo/videos/v_rBMQFpHspmo.mp4 filter=lfs diff=lfs merge=lfs -text

QWENVL-PRIVATE/.gitignore

@@ -0,0 +1,17 @@
+# python generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# venv
+.venv
+/exclude/
+.tmp.driveupload
+/excluded/
+/output/
+.idea/
+logs/
+lmms-eval-0.5.0/

QWENVL-PRIVATE/LICENSE

@@ -0,0 +1,352 @@
+Creative Commons Attribution-NonCommercial 4.0 International
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright and
+certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+-   Considerations for licensors: Our public licenses are intended for
+    use by those authorized to give the public permission to use
+    material in ways otherwise restricted by copyright and certain other
+    rights. Our licenses are irrevocable. Licensors should read and
+    understand the terms and conditions of the license they choose
+    before applying it. Licensors should also secure all rights
+    necessary before applying our licenses so that the public can reuse
+    the material as expected. Licensors should clearly mark any material
+    not subject to the license. This includes other CC-licensed
+    material, or material used under an exception or limitation to
+    copyright. More considerations for licensors :
+    wiki.creativecommons.org/Considerations_for_licensors
+
+-   Considerations for the public: By using one of our public licenses,
+    a licensor grants the public permission to use the licensed material
+    under specified terms and conditions. If the licensor's permission
+    is not necessary for any reason–for example, because of any
+    applicable exception or limitation to copyright–then that use is not
+    regulated by the license. Our licenses grant only permissions under
+    copyright and certain other rights that a licensor has authority to
+    grant. Use of the licensed material may still be restricted for
+    other reasons, including because others have copyright or other
+    rights in the material. A licensor may make special requests, such
+    as asking that all changes be marked or described. Although not
+    required by our licenses, you are encouraged to respect those
+    requests where reasonable. More considerations for the public :
+    wiki.creativecommons.org/Considerations_for_licensees
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and conditions.
+
+-   Section 1 – Definitions.
+
+    -   a. Adapted Material means material subject to Copyright and
+        Similar Rights that is derived from or based upon the Licensed
+        Material and in which the Licensed Material is translated,
+        altered, arranged, transformed, or otherwise modified in a
+        manner requiring permission under the Copyright and Similar
+        Rights held by the Licensor. For purposes of this Public
+        License, where the Licensed Material is a musical work,
+        performance, or sound recording, Adapted Material is always
+        produced where the Licensed Material is synched in timed
+        relation with a moving image.
+    -   b. Adapter's License means the license You apply to Your
+        Copyright and Similar Rights in Your contributions to Adapted
+        Material in accordance with the terms and conditions of this
+        Public License.
+    -   c. Copyright and Similar Rights means copyright and/or similar
+        rights closely related to copyright including, without
+        limitation, performance, broadcast, sound recording, and Sui
+        Generis Database Rights, without regard to how the rights are
+        labeled or categorized. For purposes of this Public License, the
+        rights specified in Section 2(b)(1)-(2) are not Copyright and
+        Similar Rights.
+    -   d. Effective Technological Measures means those measures that,
+        in the absence of proper authority, may not be circumvented
+        under laws fulfilling obligations under Article 11 of the WIPO
+        Copyright Treaty adopted on December 20, 1996, and/or similar
+        international agreements.
+    -   e. Exceptions and Limitations means fair use, fair dealing,
+        and/or any other exception or limitation to Copyright and
+        Similar Rights that applies to Your use of the Licensed
+        Material.
+    -   f. Licensed Material means the artistic or literary work,
+        database, or other material to which the Licensor applied this
+        Public License.
+    -   g. Licensed Rights means the rights granted to You subject to
+        the terms and conditions of this Public License, which are
+        limited to all Copyright and Similar Rights that apply to Your
+        use of the Licensed Material and that the Licensor has authority
+        to license.
+    -   h. Licensor means the individual(s) or entity(ies) granting
+        rights under this Public License.
+    -   i. NonCommercial means not primarily intended for or directed
+        towards commercial advantage or monetary compensation. For
+        purposes of this Public License, the exchange of the Licensed
+        Material for other material subject to Copyright and Similar
+        Rights by digital file-sharing or similar means is NonCommercial
+        provided there is no payment of monetary compensation in
+        connection with the exchange.
+    -   j. Share means to provide material to the public by any means or
+        process that requires permission under the Licensed Rights, such
+        as reproduction, public display, public performance,
+        distribution, dissemination, communication, or importation, and
+        to make material available to the public including in ways that
+        members of the public may access the material from a place and
+        at a time individually chosen by them.
+    -   k. Sui Generis Database Rights means rights other than copyright
+        resulting from Directive 96/9/EC of the European Parliament and
+        of the Council of 11 March 1996 on the legal protection of
+        databases, as amended and/or succeeded, as well as other
+        essentially equivalent rights anywhere in the world.
+    -   l. You means the individual or entity exercising the Licensed
+        Rights under this Public License. Your has a corresponding
+        meaning.
+
+-   Section 2 – Scope.
+
+    -   a. License grant.
+        -   1. Subject to the terms and conditions of this Public
+            License, the Licensor hereby grants You a worldwide,
+            royalty-free, non-sublicensable, non-exclusive, irrevocable
+            license to exercise the Licensed Rights in the Licensed
+            Material to:
+            -   A. reproduce and Share the Licensed Material, in whole
+                or in part, for NonCommercial purposes only; and
+            -   B. produce, reproduce, and Share Adapted Material for
+                NonCommercial purposes only.
+        -   2. Exceptions and Limitations. For the avoidance of doubt,
+            where Exceptions and Limitations apply to Your use, this
+            Public License does not apply, and You do not need to comply
+            with its terms and conditions.
+        -   3. Term. The term of this Public License is specified in
+            Section 6(a).
+        -   4. Media and formats; technical modifications allowed. The
+            Licensor authorizes You to exercise the Licensed Rights in
+            all media and formats whether now known or hereafter
+            created, and to make technical modifications necessary to do
+            so. The Licensor waives and/or agrees not to assert any
+            right or authority to forbid You from making technical
+            modifications necessary to exercise the Licensed Rights,
+            including technical modifications necessary to circumvent
+            Effective Technological Measures. For purposes of this
+            Public License, simply making modifications authorized by
+            this Section 2(a)(4) never produces Adapted Material.
+        -   5. Downstream recipients.
+            -   A. Offer from the Licensor – Licensed Material. Every
+                recipient of the Licensed Material automatically
+                receives an offer from the Licensor to exercise the
+                Licensed Rights under the terms and conditions of this
+                Public License.
+            -   B. No downstream restrictions. You may not offer or
+                impose any additional or different terms or conditions
+                on, or apply any Effective Technological Measures to,
+                the Licensed Material if doing so restricts exercise of
+                the Licensed Rights by any recipient of the Licensed
+                Material.
+        -   6. No endorsement. Nothing in this Public License
+            constitutes or may be construed as permission to assert or
+            imply that You are, or that Your use of the Licensed
+            Material is, connected with, or sponsored, endorsed, or
+            granted official status by, the Licensor or others
+            designated to receive attribution as provided in Section
+            3(a)(1)(A)(i).
+    -   b. Other rights.
+        -   1. Moral rights, such as the right of integrity, are not
+            licensed under this Public License, nor are publicity,
+            privacy, and/or other similar personality rights; however,
+            to the extent possible, the Licensor waives and/or agrees
+            not to assert any such rights held by the Licensor to the
+            limited extent necessary to allow You to exercise the
+            Licensed Rights, but not otherwise.
+        -   2. Patent and trademark rights are not licensed under this
+            Public License.
+        -   3. To the extent possible, the Licensor waives any right to
+            collect royalties from You for the exercise of the Licensed
+            Rights, whether directly or through a collecting society
+            under any voluntary or waivable statutory or compulsory
+            licensing scheme. In all other cases the Licensor expressly
+            reserves any right to collect such royalties, including when
+            the Licensed Material is used other than for NonCommercial
+            purposes.
+
+-   Section 3 – License Conditions.
+
+    Your exercise of the Licensed Rights is expressly made subject to
+    the following conditions.
+
+    -   a. Attribution.
+        -   1. If You Share the Licensed Material (including in modified
+            form), You must:
+            -   A. retain the following if it is supplied by the
+                Licensor with the Licensed Material:
+                -   i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if designated);
+                -   ii. a copyright notice;
+                -   iii. a notice that refers to this Public License;
+                -   iv. a notice that refers to the disclaimer of
+                    warranties;
+                -   v. a URI or hyperlink to the Licensed Material to
+                    the extent reasonably practicable;
+            -   B. indicate if You modified the Licensed Material and
+                retain an indication of any previous modifications; and
+            -   C. indicate the Licensed Material is licensed under this
+                Public License, and include the text of, or the URI or
+                hyperlink to, this Public License.
+        -   2. You may satisfy the conditions in Section 3(a)(1) in any
+            reasonable manner based on the medium, means, and context in
+            which You Share the Licensed Material. For example, it may
+            be reasonable to satisfy the conditions by providing a URI
+            or hyperlink to a resource that includes the required
+            information.
+        -   3. If requested by the Licensor, You must remove any of the
+            information required by Section 3(a)(1)(A) to the extent
+            reasonably practicable.
+        -   4. If You Share Adapted Material You produce, the Adapter's
+            License You apply must not prevent recipients of the Adapted
+            Material from complying with this Public License.
+
+-   Section 4 – Sui Generis Database Rights.
+
+    Where the Licensed Rights include Sui Generis Database Rights that
+    apply to Your use of the Licensed Material:
+
+    -   a. for the avoidance of doubt, Section 2(a)(1) grants You the
+        right to extract, reuse, reproduce, and Share all or a
+        substantial portion of the contents of the database for
+        NonCommercial purposes only;
+    -   b. if You include all or a substantial portion of the database
+        contents in a database in which You have Sui Generis Database
+        Rights, then the database in which You have Sui Generis Database
+        Rights (but not its individual contents) is Adapted Material;
+        and
+    -   c. You must comply with the conditions in Section 3(a) if You
+        Share all or a substantial portion of the contents of the
+        database.
+
+    For the avoidance of doubt, this Section 4 supplements and does not
+    replace Your obligations under this Public License where the
+    Licensed Rights include other Copyright and Similar Rights.
+
+-   Section 5 – Disclaimer of Warranties and Limitation of Liability.
+
+    -   a. Unless otherwise separately undertaken by the Licensor, to
+        the extent possible, the Licensor offers the Licensed Material
+        as-is and as-available, and makes no representations or
+        warranties of any kind concerning the Licensed Material, whether
+        express, implied, statutory, or other. This includes, without
+        limitation, warranties of title, merchantability, fitness for a
+        particular purpose, non-infringement, absence of latent or other
+        defects, accuracy, or the presence or absence of errors, whether
+        or not known or discoverable. Where disclaimers of warranties
+        are not allowed in full or in part, this disclaimer may not
+        apply to You.
+    -   b. To the extent possible, in no event will the Licensor be
+        liable to You on any legal theory (including, without
+        limitation, negligence) or otherwise for any direct, special,
+        indirect, incidental, consequential, punitive, exemplary, or
+        other losses, costs, expenses, or damages arising out of this
+        Public License or use of the Licensed Material, even if the
+        Licensor has been advised of the possibility of such losses,
+        costs, expenses, or damages. Where a limitation of liability is
+        not allowed in full or in part, this limitation may not apply to
+        You.
+    -   c. The disclaimer of warranties and limitation of liability
+        provided above shall be interpreted in a manner that, to the
+        extent possible, most closely approximates an absolute
+        disclaimer and waiver of all liability.
+
+-   Section 6 – Term and Termination.
+
+    -   a. This Public License applies for the term of the Copyright and
+        Similar Rights licensed here. However, if You fail to comply
+        with this Public License, then Your rights under this Public
+        License terminate automatically.
+    -   b. Where Your right to use the Licensed Material has terminated
+        under Section 6(a), it reinstates:
+
+        -   1. automatically as of the date the violation is cured,
+            provided it is cured within 30 days of Your discovery of the
+            violation; or
+        -   2. upon express reinstatement by the Licensor.
+
+        For the avoidance of doubt, this Section 6(b) does not affect
+        any right the Licensor may have to seek remedies for Your
+        violations of this Public License.
+
+    -   c. For the avoidance of doubt, the Licensor may also offer the
+        Licensed Material under separate terms or conditions or stop
+        distributing the Licensed Material at any time; however, doing
+        so will not terminate this Public License.
+    -   d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+        License.
+
+-   Section 7 – Other Terms and Conditions.
+
+    -   a. The Licensor shall not be bound by any additional or
+        different terms or conditions communicated by You unless
+        expressly agreed.
+    -   b. Any arrangements, understandings, or agreements regarding the
+        Licensed Material not stated herein are separate from and
+        independent of the terms and conditions of this Public License.
+
+-   Section 8 – Interpretation.
+
+    -   a. For the avoidance of doubt, this Public License does not, and
+        shall not be interpreted to, reduce, limit, restrict, or impose
+        conditions on any use of the Licensed Material that could
+        lawfully be made without permission under this Public License.
+    -   b. To the extent possible, if any provision of this Public
+        License is deemed unenforceable, it shall be automatically
+        reformed to the minimum extent necessary to make it enforceable.
+        If the provision cannot be reformed, it shall be severed from
+        this Public License without affecting the enforceability of the
+        remaining terms and conditions.
+    -   c. No term or condition of this Public License will be waived
+        and no failure to comply consented to unless expressly agreed to
+        by the Licensor.
+    -   d. Nothing in this Public License constitutes or may be
+        interpreted as a limitation upon, or waiver of, any privileges
+        and immunities that apply to the Licensor or You, including from
+        the legal processes of any jurisdiction or authority.
+
+Creative Commons is not a party to its public licenses. Notwithstanding,
+Creative Commons may elect to apply one of its public licenses to
+material it publishes and in those instances will be considered the
+"Licensor." The text of the Creative Commons public licenses is
+dedicated to the public domain under the CC0 Public Domain Dedication.
+Except for the limited purpose of indicating that material is shared
+under a Creative Commons public license or as otherwise permitted by the
+Creative Commons policies published at creativecommons.org/policies,
+Creative Commons does not authorize the use of the trademark "Creative
+Commons" or any other trademark or logo of Creative Commons without its
+prior written consent including, without limitation, in connection with
+any unauthorized modifications to any of its public licenses or any
+other arrangements, understandings, or agreements concerning use of
+licensed material. For the avoidance of doubt, this paragraph does not
+form part of the public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

QWENVL-PRIVATE/README.md

@@ -0,0 +1,1015 @@
+# Qwen2.5-VL
+
+
+<p align="center">
+    <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-VL/qwen2.5vl_logo.png" width="400"/>
+<p>
+
+<p align="center">
+        💜 <a href="https://chat.qwenlm.ai/"><b>Qwen Chat</b></a>&nbsp&nbsp | &nbsp&nbsp🤗 <a href="https://huggingface.co/collections/Qwen/qwen25-vl-6795ffac22b334a837c0f9a5">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤖 <a href="https://modelscope.cn/organization/qwen">ModelScope</a>&nbsp&nbsp | &nbsp&nbsp📑 <a href="https://qwenlm.github.io/blog/qwen2.5-vl/">Blog</a>&nbsp&nbsp | &nbsp&nbsp📚 <a href="https://github.com/QwenLM/Qwen2.5-VL/tree/main/cookbooks">Cookbooks</a>&nbsp&nbsp | &nbsp&nbsp📑 <a href="https://arxiv.org/abs/2502.13923">Paper</a>&nbsp&nbsp
+<br>
+🖥️ <a href="https://huggingface.co/spaces/Qwen/Qwen2.5-VL-72B-Instruct">Demo</a>&nbsp&nbsp | &nbsp&nbsp💬 <a href="https://github.com/QwenLM/Qwen/blob/main/assets/wechat.png">WeChat (微信)</a>&nbsp&nbsp | &nbsp&nbsp🫨 <a href="https://discord.gg/CV4E9rpNSD">Discord</a>&nbsp&nbsp | &nbsp&nbsp📑 <a href="https://help.aliyun.com/zh/model-studio/developer-reference/qwen-vl-api">API</a>&nbsp&nbsp | &nbsp&nbsp🖥️ <a href="https://gallery.pai-ml.com/#/preview/deepLearning/cv/qwen2.5-vl">PAI-DSW</a>
+</p>
+
+
+
+
+## Introduction
+In the past five months since Qwen2-VL's release, numerous developers have built new models on the Qwen2-VL vision-language models, providing us with valuable feedback. During this period, we focused on building more useful vision-language models. Today, we are excited to introduce the latest addition to the Qwen family: Qwen2.5-VL.
+
+
+#### Key Enhancements:
+
+* **Powerful Document Parsing Capabilities**: Upgrade text recognition to omnidocument parsing, excelling in processing multi-scene, multilingual, and various built-in (handwriting, tables, charts, chemical formulas, and music sheets) documents.
+
+* **Precise Object Grounding Across Formats**: Unlock improved accuracy in detecting, pointing, and counting objects, accommodating absolute coordinate and JSON formats for advanced spatial reasoning.
+
+* **Ultra-long Video Understanding and Fine-grained Video Grounding**: Extend native dynamic resolution to the temporal dimension, enhancing the ability to understand videos lasting hours while extracting event segments in seconds.
+
+* **Enhanced Agent Functionality for Computer and Mobile Devices**: Leverage advanced grounding, reasoning, and decision-making abilities, boosting the model with superior agent functionality on smartphones and computers.
+
+
+#### Model Architecture Updates:
+
+* **Dynamic Resolution and Frame Rate Training for Video Understanding**:
+
+We extend dynamic resolution to the temporal dimension by adopting dynamic FPS sampling, enabling the model to comprehend videos at various sampling rates. Accordingly, we update mRoPE in the time dimension with IDs and absolute time alignment, enabling the model to learn temporal sequence and speed, and ultimately acquire the ability to pinpoint specific moments.
+
+<p align="center">
+    <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-VL/qwen2.5vl_arc.jpeg" width="80%"/>
+<p>
+
+
+* **Streamlined and Efficient Vision Encoder**
+
+We enhance both training and inference speeds by strategically implementing window attention into the ViT. The ViT architecture is further optimized with SwiGLU and RMSNorm, aligning it with the structure of the Qwen2.5 LLM.
+
+
+## News
+* 2025.04.08: We provide the [code](https://github.com/QwenLM/Qwen2.5-VL/tree/main/qwen-vl-finetune) for fine-tuning Qwen2-VL and Qwen2.5-VL.
+* 2025.03.25: We have released the [Qwen2.5-VL-32B](https://huggingface.co/Qwen/Qwen2.5-VL-32B-Instruct). It is smarter and its responses align more closely with human preferences. For more details, please check our [blog](https://qwenlm.github.io/blog/qwen2.5-vl-32b/)!
+* 2025.02.20: we have released the [Qwen2.5-VL Technical Report](https://arxiv.org/abs/2502.13923). Alongside the report, we have also released AWQ-quantized models for Qwen2.5-VL in three different sizes: [3B](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct-AWQ), [7B](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct-AWQ) , and [72B](https://huggingface.co/Qwen/Qwen2.5-VL-72B-Instruct-AWQ) parameters.
+* 2025.01.28: We have released the [Qwen2.5-VL series](https://huggingface.co/Qwen). For more details, please check our [blog](https://qwenlm.github.io/blog/qwen2.5-vl/)!
+* 2024.12.25: We have released the [QvQ-72B-Preview](https://huggingface.co/Qwen/QVQ-72B-Preview). QvQ-72B-Preview is an experimental research model, focusing on enhancing visual reasoning capabilities. For more details, please check our [blog](https://qwenlm.github.io/blog/qvq-72b-preview/)!
+* 2024.09.19: The instruction-tuned [Qwen2-VL-72B model](https://huggingface.co/Qwen/Qwen2-VL-72B-Instruct) and its quantized version [[AWQ](https://huggingface.co/Qwen/Qwen2-VL-72B-Instruct-AWQ), [GPTQ-Int4](https://huggingface.co/Qwen/Qwen2-VL-72B-Instruct-GPTQ-Int4), [GPTQ-Int8](https://huggingface.co/Qwen/Qwen2-VL-72B-Instruct-GPTQ-Int8)] are now available. We have also released the [Qwen2-VL paper](https://arxiv.org/pdf/2409.12191) simultaneously.
+* 2024.08.30: We have released the [Qwen2-VL series](https://huggingface.co/collections/Qwen/qwen2-vl-66cee7455501d7126940800d). The 2B and 7B models are now available, and the 72B model for open source is coming soon. For more details, please check our [blog](https://qwenlm.github.io/blog/qwen2-vl/)!
+
+
+## Performance
+
+
+
+| Dataset            | Qwen2.5-VL-3B<br><sup>([🤗](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct)[🤖](https://modelscope.cn/models/qwen/Qwen2.5-VL-3B-Instruct))     | Qwen2.5-VL-7B<br><sup>([🤗](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct)[🤖](https://modelscope.cn/models/qwen/Qwen2.5-VL-7B-Instruct))    | Qwen2.5-VL-32B<br><sup>([🤗](https://huggingface.co/Qwen/Qwen2.5-VL-32B-Instruct)[🤖](https://modelscope.cn/models/qwen/Qwen2.5-VL-32B-Instruct))    | Qwen2.5-VL-72B<br><sup>([🤗](https://huggingface.co/Qwen/Qwen2.5-VL-72B-Instruct)[🤖](https://modelscope.cn/models/qwen/Qwen2.5-VL-72B-Instruct)) | Gemini-2 Flash | GPT-4o | Claude3.5 Sonnet | Qwen2-VL 72B | 
+|--------------------|--------|--------|--------|----------------|----------------|--------|------------------|--------------|
+| MMMU               | 53.1  | 58.6 | 70.0  | **70.2**           | **70.7**           | 70.3   | 70.4             | 64.5         | 
+| MMMU Pro           | 31.6  | 38.3 |49.5|**51.1**           | **57**             | 54.5   | 54.7             | 46.2         | 
+| DocVQA             | 93.9 | 95.7 | 94.8 |**96.4**           | 92.1           | 91.1   | 95.2             | **96.5**         | 
+| InfoVQA            | 77.1 | 82.6 | 83.4| **87.3**          | 77.8           | 80.7   | 74.3             | 84.5         | 
+| CC-OCR             | 74.5 |  77.8 | 77.1 |**79.8**          | 73.0           | 66.6   | 62.7             | 68.7         | 
+| OCRBenchV2         | 54.3/52.1 | 56.3/57.2 | 57.2/59.1| **61.5/63.7**      | -              | 46.5/32.3 | 45.2/39.6       | 47.8/46.1    | 
+| MegaBench          | 28.9 |36.8 | - |**51.3**          | **55.2**           | 54.2   | 52.1             | 46.8         | 
+| MMStar             | 55.8|63.9| 69.5 |**70.8**           | 69.4           | 64.7   | 65.1             | 68.3         | 
+| MMBench1.1         |  81.5  | 84.3 | 84.6|**88.0**           | 83.0           | 82.1   | 83.4             | 86.6         | 
+| MathVista          | 62.3  | 68.2  | 74.7 |**74.8**           | 73.1           | 63.8   | 65.4             | 70.5         | 
+| MathVision         | 21.2  | 25.1 | **38.4** |38.1           | **41.3**           | 30.4   | 38.3             | 25.9         | 
+| VideoMME           | 61.5/67.6 | 65.1/71.6 | 70.5/77.9 |**73.3/79.1**      | -/-            | 71.9/77.2 | 60/62.9         | 71.2/77.8    | 
+| MMBench-Video      | 1.63           | 1.79    | 1.93 |**2.02**           | -              | 1.68   | 1.38             | 1.7          | 
+| LVBench            | 43.3        | 45.3 | **49.0** | 47.3           | -              | 30.8   | -                | -            | 
+| CharadesSTA        | 38.8        | 43.6    |   **54.2** | 50.9           | -              | 35.7   | -                | -            | 
+| AITZ               | 76.9|81.9 | 83.1 |**83.2**           | -              | 35.3   | -                | -            | 
+| Android Control    | 63.7/90.8 | 60.1/91.4 | 69.6/93.3 |**67.36/93.7**     | -              | -      | -                | 66.4/84.4    | 
+| ScreenSpot         | 55.5|84.7| **88.5** | 87.1          | 84.0           | 18.1   | 83.0             | -            | 
+| ScreenSpot Pro     | 23.9|29.0 |  39.4 |**43.6**           | -              | -      | 17.1             | -            | 
+| AndroidWorld       | -  | -  | 22.0 |**35**             | -              | 34.5(SoM) | 27.9            | -            | 
+| OSWorld            | -  | -  | 5.92 |**8.83**           | -              | 5.03   | **14.9**             | -            | 
+
+
+
+
+
+## Quickstart
+
+Below, we provide simple examples to show how to use Qwen2.5-VL with 🤖 ModelScope and 🤗 Transformers.
+
+The code of Qwen2.5-VL has been in the latest Hugging face transformers and we advise you to build from source with command:
+```
+pip install transformers==4.51.3 accelerate
+```
+or you might encounter the following error:
+```
+KeyError: 'qwen2_5_vl'
+```
+
+
+We offer a toolkit to help you handle various types of visual input more conveniently, as if you were using an API. This includes base64, URLs, and interleaved images and videos. You can install it using the following command:
+
+```bash
+# It's highly recommended to use `[decord]` feature for faster video loading.
+pip install qwen-vl-utils[decord]
+```
+
+Currently, `qwen-vl-utils` supports three video decoding backends: `torchvision`, `decord`, and `torchcodec`. While `decord` and `torchcodec` generally offer significantly faster decoding speeds compared to `torchvision`, we recommend using `torchcodec`. This is because `decord` has known issues, such as decoding hangs, and its project is no longer actively maintained.
+
+- For `decord`, if you are not using Linux, you might not be able to install `decord` from PyPI. In that case, you can use `pip install qwen-vl-utils` which will fall back to using torchvision for video processing. However, you can still [install decord from source](https://github.com/dmlc/decord?tab=readme-ov-file#install-from-source) to get decord used when loading video.
+
+- To use `torchcodec` as the backend for video decoding, follow the installation instructions provided in the official [torchcodec repository](https://github.com/pytorch/torchcodec/tree/main?tab=readme-ov-file#installing-torchcodec) and install it manually. Note that `torchcodec` depends on FFmpeg for decoding functionality.
+
+## Cookbooks
+
+We are preparing [cookbooks](https://github.com/QwenLM/Qwen2.5-VL/tree/main/cookbooks) for many capabilities, including recognition, localization, document parsing, video understanding, key information extraction, and more. Welcome to learn more!
+
+| Cookbook | Description | Open |
+| -------- | ----------- | ---- |
+| [Universal Recognition](https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/universal_recognition.ipynb) | Not only identify animals, plants, people, and scenic spots but also recognize various objects such as cars and merchandise. | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QwenLM/Qwen2.5-VL/blob/main/cookbooks/universal_recognition.ipynb) |
+| [Powerful Document Parsing Capabilities](https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/document_parsing.ipynb) | The parsing of documents has reached a higher level, including not only text but also layout position information and our Qwen HTML format. | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QwenLM/Qwen2.5-VL/blob/main/cookbooks/document_parsing.ipynb) |
+| [Precise Object Grounding Across Formats](https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/spatial_understanding.ipynb) | Using absolute position coordinates, it supports both boxes and points, allowing for diverse combinations of positioning and labeling tasks. | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QwenLM/Qwen2.5-VL/blob/main/cookbooks/spatial_understanding.ipynb) |
+| [General OCR and Key Information Extraction](https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/ocr.ipynb) | Stronger text recognition capabilities in natural scenes and multiple languages, supporting diverse key information extraction needs. | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QwenLM/Qwen2.5-VL/blob/main/cookbooks/ocr.ipynb) |
+| [Video Understanding](https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/video_understanding.ipynb) | Better video OCR, long video understanding, and video grounding. | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QwenLM/Qwen2.5-VL/blob/main/cookbooks/video_understanding.ipynb) |
+| [Mobile Agent](https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/mobile_agent.ipynb) | Locate and think for mobile phone control. | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QwenLM/Qwen2.5-VL/blob/main/cookbooks/mobile_agent.ipynb) |
+| [Computer-Use Agent](https://github.com/QwenLM/Qwen2.5-VL/blob/main/cookbooks/computer_use.ipynb) | Locate and think for controlling computers and Web. | [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QwenLM/Qwen2.5-VL/blob/main/cookbooks/computer_use.ipynb) |
+
+
+### Using 🤗  Transformers to Chat
+
+Here we show a code snippet to show you how to use the chat model with `transformers` and `qwen_vl_utils`:
+
+```python
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+# default: Load the model on the available device(s)
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
+)
+
+# We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
+# model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+#     "Qwen/Qwen2.5-VL-7B-Instruct",
+#     torch_dtype=torch.bfloat16,
+#     attn_implementation="flash_attention_2",
+#     device_map="auto",
+# )
+
+# default processor
+processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
+
+# The default range for the number of visual tokens per image in the model is 4-16384.
+# You can set min_pixels and max_pixels according to your needs, such as a token range of 256-1280, to balance performance and cost.
+# min_pixels = 256*28*28
+# max_pixels = 1280*28*28
+# processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
+
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+            },
+            {"type": "text", "text": "Describe this image."},
+        ],
+    }
+]
+
+# Preparation for inference
+text = processor.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+image_inputs, video_inputs = process_vision_info(messages)
+inputs = processor(
+    text=[text],
+    images=image_inputs,
+    videos=video_inputs,
+    padding=True,
+    return_tensors="pt",
+)
+inputs = inputs.to(model.device)
+
+# Inference: Generation of the output
+generated_ids = model.generate(**inputs, max_new_tokens=128)
+generated_ids_trimmed = [
+    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+]
+output_text = processor.batch_decode(
+    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+)
+print(output_text)
+```
+
+<details>
+<summary>Minimum VRAM requirements</summary>
+
+| Precision | Qwen2.5-VL-3B | Qwen2.5-VL-7B | Qwen2.5-VL-72B |
+|-----------|------------| --------- | -------- |
+| FP32      | 11.5 GB    | 26.34 GB  | 266.21 GB |
+| BF16      | 5.75 GB    | 13.17 GB  | 133.11 GB |
+| INT8      | 2.87 GB    | 6.59 GB   | 66.5 GB |
+| INT4      | 1.44 GB    | 3.29 GB   | 33.28 GB |
+
+Note: The table above presents the theoretical minimum video memory requirements for inference with `transformers`; however, in practice, the actual memory usage is typically at least 1.2 times higher. For more information, see the linked resource [here](https://huggingface.co/docs/accelerate/main/en/usage_guides/model_size_estimator).
+</details>
+
+
+<details>
+<summary>Multi image inference</summary>
+
+```python
+# Messages containing multiple images and a text query
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "file:///path/to/image1.jpg"},
+            {"type": "image", "image": "file:///path/to/image2.jpg"},
+            {"type": "text", "text": "Identify the similarities between these images."},
+        ],
+    }
+]
+
+# Preparation for inference
+text = processor.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+image_inputs, video_inputs = process_vision_info(messages)
+inputs = processor(
+    text=[text],
+    images=image_inputs,
+    videos=video_inputs,
+    padding=True,
+    return_tensors="pt",
+)
+inputs = inputs.to("cuda")
+
+# Inference
+generated_ids = model.generate(**inputs, max_new_tokens=128)
+generated_ids_trimmed = [
+    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+]
+output_text = processor.batch_decode(
+    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+)
+print(output_text)
+```
+</details>
+
+<details>
+<summary>Video inference</summary>
+
+```python
+# Messages containing a images list as a video and a text query
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "video",
+                "video": [
+                    "file:///path/to/frame1.jpg",
+                    "file:///path/to/frame2.jpg",
+                    "file:///path/to/frame3.jpg",
+                    "file:///path/to/frame4.jpg",
+                ],
+            },
+            {"type": "text", "text": "Describe this video."},
+        ],
+    }
+]
+
+# Messages containing a local video path and a text query
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "video",
+                "video": "file:///path/to/video1.mp4",
+                "max_pixels": 360 * 420,
+                "fps": 1.0,
+            },
+            {"type": "text", "text": "Describe this video."},
+        ],
+    }
+]
+
+# Messages containing a video url and a text query
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "video",
+                "video": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-VL/space_woaudio.mp4",
+                "min_pixels": 4 * 28 * 28,
+                "max_pixels": 256 * 28 * 28,
+                "total_pixels": 20480 * 28 * 28,
+            },
+            {"type": "text", "text": "Describe this video."},
+        ],
+    }
+]
+
+# Preparation for inference
+text = processor.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+image_inputs, video_inputs, video_kwargs = process_vision_info(messages, return_video_kwargs=True)
+inputs = processor(
+    text=[text],
+    images=image_inputs,
+    videos=video_inputs,
+    fps=fps,
+    padding=True,
+    return_tensors="pt",
+    **video_kwargs,
+)
+inputs = inputs.to("cuda")
+
+# Inference
+generated_ids = model.generate(**inputs, max_new_tokens=128)
+generated_ids_trimmed = [
+    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+]
+output_text = processor.batch_decode(
+    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+)
+print(output_text)
+```
+
+#### Video URL compatibility
+
+Video URL compatibility is primarily determined by the version of the third-party library being used. For more details, refer to the table below. If you prefer not to use the default backend, you can switch it by setting `FORCE_QWENVL_VIDEO_READER` to `torchvision`, `decord`, or `torchcodec`.
+
+| Backend     | HTTP | HTTPS |
+|-------------|------|-------|
+| torchvision >= 0.19.0 | ✅  | ✅   |
+| torchvision < 0.19.0  | ❌  | ❌   |
+| decord      | ✅  | ❌   |
+| torchcodec  | ✅  | ✅   |
+
+#### Configuration for adjusting video resolution
+
+We recommend setting appropriate values for the `min_pixels` and `max_pixels` parameters based on available GPU memory and the specific application scenario to restrict the resolution of individual frames in the video. Alternatively, you can use the `total_pixels` parameter to limit the total number of tokens in the video (it is recommended to set this value below 24576 * 28 * 28 to avoid excessively long input sequences). For more details on parameter usage and processing logic, please refer to the `fetch_video` function in `qwen_vl_utils/vision_process.py`.
+
+</details>
+
+<details>
+<summary>Batch inference</summary>
+
+```python
+# Sample messages for batch inference
+messages1 = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "file:///path/to/image1.jpg"},
+            {"type": "image", "image": "file:///path/to/image2.jpg"},
+            {"type": "text", "text": "What are the common elements in these pictures?"},
+        ],
+    }
+]
+messages2 = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Who are you?"},
+]
+# Combine messages for batch processing
+messages = [messages1, messages2]
+
+# Preparation for batch inference
+texts = [
+    processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+    for msg in messages
+]
+image_inputs, video_inputs = process_vision_info(messages)
+inputs = processor(
+    text=texts,
+    images=image_inputs,
+    videos=video_inputs,
+    padding=True,
+    return_tensors="pt",
+)
+inputs = inputs.to("cuda")
+
+# Batch Inference
+generated_ids = model.generate(**inputs, max_new_tokens=128)
+generated_ids_trimmed = [
+    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+]
+output_texts = processor.batch_decode(
+    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+)
+print(output_texts)
+```
+</details>
+
+### 🤖 ModelScope
+We strongly advise users especially those in mainland China to use ModelScope. `snapshot_download` can help you solve issues concerning downloading checkpoints.
+
+### More Usage Tips
+
+For input images, we support local files, base64, and URLs. For videos, we currently only support local files.
+
+```python
+# You can directly insert a local file path, a URL, or a base64-encoded image into the position where you want in the text.
+## Local file path
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "file:///path/to/your/image.jpg"},
+            {"type": "text", "text": "Describe this image."},
+        ],
+    }
+]
+## Image URL
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "http://path/to/your/image.jpg"},
+            {"type": "text", "text": "Describe this image."},
+        ],
+    }
+]
+## Base64 encoded image
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "data:image;base64,/9j/..."},
+            {"type": "text", "text": "Describe this image."},
+        ],
+    }
+]
+```
+#### Image Resolution for performance boost
+
+The model supports a wide range of resolution inputs. By default, it uses the native resolution for input, but higher resolutions can enhance performance at the cost of more computation. Users can set the minimum and maximum number of pixels to achieve an optimal configuration for their needs, such as a token count range of 256-1280, to balance speed and memory usage.
+
+```python
+min_pixels = 256 * 28 * 28
+max_pixels = 1280 * 28 * 28
+processor = AutoProcessor.from_pretrained(
+    "Qwen/Qwen2.5-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels
+)
+```
+
+Besides, We provide two methods for fine-grained control over the image size input to the model:
+
+1. Specify exact dimensions: Directly set `resized_height` and `resized_width`. These values will be rounded to the nearest multiple of 28.
+
+2. Define min_pixels and max_pixels: Images will be resized to maintain their aspect ratio within the range of min_pixels and max_pixels.
+
+```python
+# resized_height and resized_width
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "image": "file:///path/to/your/image.jpg",
+                "resized_height": 280,
+                "resized_width": 420,
+            },
+            {"type": "text", "text": "Describe this image."},
+        ],
+    }
+]
+# min_pixels and max_pixels
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "image": "file:///path/to/your/image.jpg",
+                "min_pixels": 50176,
+                "max_pixels": 50176,
+            },
+            {"type": "text", "text": "Describe this image."},
+        ],
+    }
+]
+```
+
+#### Add ids for Multiple Image Inputs
+By default, images and video content are directly included in the conversation. When handling multiple images, it's helpful to add labels to the images and videos for better reference. Users can control this behavior with the following settings:
+<details>
+<summary>Add vision ids</summary>
+
+```python
+conversation = [
+    {
+        "role": "user",
+        "content": [{"type": "image"}, {"type": "text", "text": "Hello, how are you?"}],
+    },
+    {
+        "role": "assistant",
+        "content": "I'm doing well, thank you for asking. How can I assist you today?",
+    },
+    {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "Can you describe these images and video?"},
+            {"type": "image"},
+            {"type": "image"},
+            {"type": "video"},
+            {"type": "text", "text": "These are from my vacation."},
+        ],
+    },
+    {
+        "role": "assistant",
+        "content": "I'd be happy to describe the images and video for you. Could you please provide more context about your vacation?",
+    },
+    {
+        "role": "user",
+        "content": "It was a trip to the mountains. Can you see the details in the images and video?",
+    },
+]
+
+# default:
+prompt_without_id = processor.apply_chat_template(
+    conversation, add_generation_prompt=True
+)
+# Excepted output: '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Hello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing well, thank you for asking. How can I assist you today?<|im_end|>\n<|im_start|>user\nCan you describe these images and video?<|vision_start|><|image_pad|><|vision_end|><|vision_start|><|image_pad|><|vision_end|><|vision_start|><|video_pad|><|vision_end|>These are from my vacation.<|im_end|>\n<|im_start|>assistant\nI'd be happy to describe the images and video for you. Could you please provide more context about your vacation?<|im_end|>\n<|im_start|>user\nIt was a trip to the mountains. Can you see the details in the images and video?<|im_end|>\n<|im_start|>assistant\n'
+
+
+# add ids
+prompt_with_id = processor.apply_chat_template(
+    conversation, add_generation_prompt=True, add_vision_id=True
+)
+# Excepted output: '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nPicture 1: <|vision_start|><|image_pad|><|vision_end|>Hello, how are you?<|im_end|>\n<|im_start|>assistant\nI'm doing well, thank you for asking. How can I assist you today?<|im_end|>\n<|im_start|>user\nCan you describe these images and video?Picture 2: <|vision_start|><|image_pad|><|vision_end|>Picture 3: <|vision_start|><|image_pad|><|vision_end|>Video 1: <|vision_start|><|video_pad|><|vision_end|>These are from my vacation.<|im_end|>\n<|im_start|>assistant\nI'd be happy to describe the images and video for you. Could you please provide more context about your vacation?<|im_end|>\n<|im_start|>user\nIt was a trip to the mountains. Can you see the details in the images and video?<|im_end|>\n<|im_start|>assistant\n'
+```
+</details>
+
+#### Flash-Attention 2 to speed up generation
+
+First, make sure to install the latest version of Flash Attention 2:
+
+```bash
+pip install -U flash-attn --no-build-isolation
+```
+
+Also, you should have a hardware that is compatible with Flash-Attention 2. Read more about it in the official documentation of the [flash attention repository](https://github.com/Dao-AILab/flash-attention). FlashAttention-2 can only be used when a model is loaded in `torch.float16` or `torch.bfloat16`.
+
+To load and run a model using Flash Attention-2, simply add `attn_implementation="flash_attention_2"` when loading the model as follows:
+
+```python
+from transformers import Qwen2_5_VLForConditionalGeneration
+
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    "Qwen/Qwen2.5-VL-7B-Instruct", 
+    torch_dtype=torch.bfloat16, 
+    attn_implementation="flash_attention_2",
+)
+```
+
+### Processing Long Texts
+
+The current `config.json` is set for context length up to 32,768 tokens.
+To handle extensive inputs exceeding 32,768 tokens, we utilize [YaRN](https://arxiv.org/abs/2309.00071), a technique for enhancing model length extrapolation, ensuring optimal performance on lengthy texts.
+
+For supported frameworks, you could add the following to `config.json` to enable YaRN:
+
+```
+{
+	...,
+    "type": "yarn",
+    "mrope_section": [
+        16,
+        24,
+        24
+    ],
+    "factor": 4,
+    "original_max_position_embeddings": 32768
+}
+```
+
+However, it should be noted that this method has a significant impact on the performance of temporal and spatial localization tasks, and is therefore not recommended for use.
+
+At the same time, for long video inputs, since MRoPE itself is more economical with ids, the max_position_embeddings can be directly modified to a larger value, such as 64k.
+
+
+### Try Qwen2.5-VL-72B with API!
+
+To explore Qwen2.5-VL-72B, a more fascinating multimodal model, we encourage you to test our cutting-edge API service. Let's start the exciting journey right now!
+
+#### Installation
+```bash
+pip install dashscope
+```
+
+#### Examples
+```python
+import dashscope
+
+
+dashscope.api_key = "your_api_key"
+
+messages = [{
+    'role': 'user',
+    'content': [
+        {
+            'image': "https://dashscope.oss-cn-beijing.aliyuncs.com/images/dog_and_girl.jpeg"
+        },
+        {
+            'text': 'What are in the image?'
+        },
+    ]
+}]
+
+response = dashscope.MultiModalConversation.call(model='qwen2.5-vl-72b-instruct', messages=messages)
+print(response)
+```
+
+For more usage, please refer to the tutorial at [aliyun](https://help.aliyun.com/zh/model-studio/developer-reference/qwen-vl-api).
+
+
+
+
+## Demo
+### Web UI Example
+
+In this section, we provide instructions for users to build a web-based user interface (UI) demo. This UI demo allows users to interact with a predefined model or application through a web browser. Follow the steps below to get started.
+
+#### Installation
+
+Before you begin, ensure that you have the required dependencies installed on your system. You can install them by running the following command:
+
+```bash
+pip install -r requirements_web_demo.txt
+```
+
+#### Running the Demo with FlashAttention-2
+
+Once the required packages are installed, you can launch the web demo using the following command. This command will start a web server and provide you with a link to access the UI in your web browser.
+
+**Recommended**: For enhanced performance and efficiency, especially in multi-image and video processing scenarios, we strongly recommend using [FlashAttention-2](https://github.com/Dao-AILab/flash-attention). FlashAttention-2 provides significant improvements in memory usage and speed, making it ideal for handling large-scale models and data processing.
+
+To enable FlashAttention-2, use the following command:
+
+```bash
+python web_demo_mm.py --flash-attn2
+```
+
+This will load the model with FlashAttention-2 enabled.
+
+**Default Usage**: If you prefer to run the demo without FlashAttention-2 or if you do not specify the `--flash-attn2` option, the demo will load the model using the standard attention implementation:
+
+```bash
+python web_demo_mm.py
+```
+
+After running the command, you’ll see a link generated in the terminal similar to this:
+
+```
+Running on local: http://127.0.0.1:7860/
+```
+
+Copy this link and paste it into your browser to access the web UI, where you can interact with the model by inputting text, uploading images, or using any other provided functionalities.
+
+##### Running the Streaming Video Chat Demo
+An experimental streaming video chat demo is also available in the ``web_demo_streaming`` directory.
+
+To run the streaming video chat demo, use the following command:
+
+```bash
+cd web_demo_streaming/
+python app.py --flash-attn2
+```
+
+If you prefer to run the demo without FlashAttention-2, use the following command:
+```bash
+cd web_demo_streaming/
+python app.py
+```
+
+This demo supports webcam/screen capture as its video input source. To support screen capture video input, we use code snippet from the following hugginface space: [gstaff/gradio-screen-recorder](https://huggingface.co/spaces/gstaff/gradio-screen-recorder/tree/main).
+
+
+## Deployment
+
+We recommend using vLLM for fast Qwen2.5-VL deployment and inference. You need to install `vllm>0.7.2` to enable Qwen2.5-VL support. You can also use our [official docker image](#-docker).
+
+You can also check [vLLM official documentation](https://docs.vllm.ai/en/latest/serving/multimodal_inputs.html) for more details about online serving and offline inference.
+
+### Installation
+```bash
+pip install git+https://github.com/huggingface/transformers@f3f6c86582611976e72be054675e2bf0abb5f775
+pip install accelerate
+pip install qwen-vl-utils
+pip install 'vllm>0.7.2'
+
+```
+### Start an OpenAI API Service
+
+Run the command below to start an OpenAI-compatible API service:
+
+```bash
+vllm serve Qwen/Qwen2.5-VL-7B-Instruct --port 8000 --host 0.0.0.0 --dtype bfloat16 --limit-mm-per-prompt image=5,video=5
+```
+
+Then you can use the chat API as below (via curl or Python API):
+
+```bash
+curl http://localhost:8000/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d '{
+    "model": "Qwen/Qwen2.5-VL-7B-Instruct",
+    "messages": [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": [
+        {"type": "image_url", "image_url": {"url": "https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png"}},
+        {"type": "text", "text": "What is the text in the illustrate?"}
+    ]}
+    ]
+    }'
+```
+
+```python
+from openai import OpenAI
+
+# Set OpenAI's API key and API base to use vLLM's API server.
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:8000/v1"
+
+client = OpenAI(
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+chat_response = client.chat.completions.create(
+    model="Qwen/Qwen2.5-VL-7B-Instruct",
+    messages=[
+        {"role": "system", "content": "You are a helpful assistant."},
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": "https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png"
+                    },
+                },
+                {"type": "text", "text": "What is the text in the illustrate?"},
+            ],
+        },
+    ],
+)
+print("Chat response:", chat_response)
+```
+
+You can also upload base64-encoded local images (see [OpenAI API protocol document](https://platform.openai.com/docs/guides/vision/uploading-base-64-encoded-images) for more details):
+```python
+import base64
+from openai import OpenAI
+# Set OpenAI's API key and API base to use vLLM's API server.
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:8000/v1"
+client = OpenAI(
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+image_path = "/path/to/local/image.png"
+with open(image_path, "rb") as f:
+    encoded_image = base64.b64encode(f.read())
+encoded_image_text = encoded_image.decode("utf-8")
+base64_qwen = f"data:image;base64,{encoded_image_text}"
+chat_response = client.chat.completions.create(
+    model="Qwen/Qwen2.5-VL-7B-Instruct",
+    messages=[
+        {"role": "system", "content": "You are a helpful assistant."},
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": base64_qwen
+                    },
+                },
+                {"type": "text", "text": "What is the text in the illustrate?"},
+            ],
+        },
+    ],
+)
+print("Chat response:", chat_response)
+```
+
+For videos, you can use the chat API  as follows：
+```python
+import base64
+import numpy as np
+from PIL import Image
+from io import BytesIO
+from openai import OpenAI
+from qwen_vl_utils import process_vision_info
+
+
+# Set OpenAI's API key and API base to use vLLM's API server.
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:8000/v1"
+
+client = OpenAI(
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+
+video_messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": [
+        {"type": "text", "text": "请用表格总结一下视频中的商品特点"},
+        {
+            "type": "video",
+            "video": "https://duguang-labelling.oss-cn-shanghai.aliyuncs.com/qiansun/video_ocr/videos/50221078283.mp4",
+            "total_pixels": 20480 * 28 * 28, "min_pixels": 16 * 28 * 2, 
+            'fps': 3.0  # The default value is 2.0, but for demonstration purposes, we set it to 3.0.
+        }]
+    },
+]
+
+
+def prepare_message_for_vllm(content_messages):
+    """
+    The frame extraction logic for videos in `vLLM` differs from that of `qwen_vl_utils`.
+    Here, we utilize `qwen_vl_utils` to extract video frames, with the `media_typ`e of the video explicitly set to `video/jpeg`.
+    By doing so, vLLM will no longer attempt to extract frames from the input base64-encoded images.
+    """
+    vllm_messages, fps_list = [], []
+    for message in content_messages:
+        message_content_list = message["content"]
+        if not isinstance(message_content_list, list):
+            vllm_messages.append(message)
+            continue
+
+        new_content_list = []
+        for part_message in message_content_list:
+            if 'video' in part_message:
+                video_message = [{'content': [part_message]}]
+                image_inputs, video_inputs, video_kwargs = process_vision_info(video_message, return_video_kwargs=True)
+                assert video_inputs is not None, "video_inputs should not be None"
+                video_input = (video_inputs.pop()).permute(0, 2, 3, 1).numpy().astype(np.uint8)
+                fps_list.extend(video_kwargs.get('fps', []))
+
+                # encode image with base64
+                base64_frames = []
+                for frame in video_input:
+                    img = Image.fromarray(frame)
+                    output_buffer = BytesIO()
+                    img.save(output_buffer, format="jpeg")
+                    byte_data = output_buffer.getvalue()
+                    base64_str = base64.b64encode(byte_data).decode("utf-8")
+                    base64_frames.append(base64_str)
+
+                part_message = {
+                    "type": "video_url",
+                    "video_url": {"url": f"data:video/jpeg;base64,{','.join(base64_frames)}"}
+                }
+            new_content_list.append(part_message)
+        message["content"] = new_content_list
+        vllm_messages.append(message)
+    return vllm_messages, {'fps': fps_list}
+
+
+video_messages, video_kwargs = prepare_message_for_vllm(video_messages)
+chat_response = client.chat.completions.create(
+    model="Qwen/Qwen2.5-VL-7B-Instruct",
+    messages=video_messages,
+    extra_body={
+        "mm_processor_kwargs": video_kwargs
+    }
+)
+print("Chat response:", chat_response)
+```
+
+### Inference Locally
+
+You can also use vLLM to inference Qwen2.5-VL locally:
+
+```python
+from transformers import AutoProcessor
+from vllm import LLM, SamplingParams
+from qwen_vl_utils import process_vision_info
+
+MODEL_PATH = "Qwen/Qwen2.5-VL-7B-Instruct"
+
+llm = LLM(
+    model=MODEL_PATH,
+    limit_mm_per_prompt={"image": 10, "video": 10},
+)
+
+sampling_params = SamplingParams(
+    temperature=0.1,
+    top_p=0.001,
+    repetition_penalty=1.05,
+    max_tokens=256,
+    stop_token_ids=[],
+)
+
+image_messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "image": "https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png",
+                "min_pixels": 224 * 224,
+                "max_pixels": 1280 * 28 * 28,
+            },
+            {"type": "text", "text": "What is the text in the illustrate?"},
+        ],
+    },
+]
+
+
+# For video input, you can pass following values instead:
+# "type": "video",
+# "video": "<video URL>",
+video_messages = [
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": [
+            {"type": "text", "text": "请用表格总结一下视频中的商品特点"},
+            {
+                "type": "video", 
+                "video": "https://duguang-labelling.oss-cn-shanghai.aliyuncs.com/qiansun/video_ocr/videos/50221078283.mp4",
+                "total_pixels": 20480 * 28 * 28, "min_pixels": 16 * 28 * 28
+            }
+        ]
+    },
+]
+
+# Here we use video messages as a demonstration
+messages = video_messages
+
+processor = AutoProcessor.from_pretrained(MODEL_PATH)
+prompt = processor.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+)
+image_inputs, video_inputs, video_kwargs = process_vision_info(messages, return_video_kwargs=True)
+
+mm_data = {}
+if image_inputs is not None:
+    mm_data["image"] = image_inputs
+if video_inputs is not None:
+    mm_data["video"] = video_inputs
+
+llm_inputs = {
+    "prompt": prompt,
+    "multi_modal_data": mm_data,
+
+    # FPS will be returned in video_kwargs
+    "mm_processor_kwargs": video_kwargs,
+}
+
+outputs = llm.generate([llm_inputs], sampling_params=sampling_params)
+generated_text = outputs[0].outputs[0].text
+
+print(generated_text)
+```
+
+
+## 🐳 Docker
+
+To simplify the deploy process, we provide docker images with pre-build environments: [qwenllm/qwenvl](https://hub.docker.com/r/qwenllm/qwenvl). You only need to install the driver and download model files to launch demos.
+
+```bash
+docker run --gpus all --ipc=host --network=host --rm --name qwen2.5 -it qwenllm/qwenvl:2.5-cu121 bash
+```
+
+## Citation
+
+If you find our paper and code useful in your research, please consider giving a star :star: and citation :pencil: :)
+
+
+
+
+```BibTeX
+
+@article{Qwen2.5-VL,
+  title={Qwen2.5-VL Technical Report},
+  author={Bai, Shuai and Chen, Keqin and Liu, Xuejing and Wang, Jialin and Ge, Wenbin and Song, Sibo and Dang, Kai and Wang, Peng and Wang, Shijie and Tang, Jun and Zhong, Humen and Zhu, Yuanzhi and Yang, Mingkun and Li, Zhaohai and Wan, Jianqiang and Wang, Pengfei and Ding, Wei and Fu, Zheren and Xu, Yiheng and Ye, Jiabo and Zhang, Xi and Xie, Tianbao and Cheng, Zesen and Zhang, Hang and Yang, Zhibo and Xu, Haiyang and Lin, Junyang},
+  journal={arXiv preprint arXiv:2502.13923},
+  year={2025}
+}
+
+@article{Qwen2-VL,
+  title={Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution},
+  author={Wang, Peng and Bai, Shuai and Tan, Sinan and Wang, Shijie and Fan, Zhihao and Bai, Jinze and Chen, Keqin and Liu, Xuejing and Wang, Jialin and Ge, Wenbin and Fan, Yang and Dang, Kai and Du, Mengfei and Ren, Xuancheng and Men, Rui and Liu, Dayiheng and Zhou, Chang and Zhou, Jingren and Lin, Junyang},
+  journal={arXiv preprint arXiv:2409.12191},
+  year={2024}
+}
+
+@article{Qwen-VL,
+  title={Qwen-VL: A Versatile Vision-Language Model for Understanding, Localization, Text Reading, and Beyond},
+  author={Bai, Jinze and Bai, Shuai and Yang, Shusheng and Wang, Shijie and Tan, Sinan and Wang, Peng and Lin, Junyang and Zhou, Chang and Zhou, Jingren},
+  journal={arXiv preprint arXiv:2308.12966},
+  year={2023}
+}
+```
+
+<br>

QWENVL-PRIVATE/analysis/aaten_debug.py

@@ -0,0 +1,393 @@
+import sys
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import AutoProcessor
+from qwen_src.qwen3_vl import Qwen3VLForConditionalGeneration
+from qwen_src.qwen3_vl.aura_attention import AuraDecoderLayer
+import json
+import os
+from tqdm import tqdm
+from PIL import Image, ImageDraw, ImageFont
+from qwen_src.mm_utils import expand2square, create_pseudo_labels
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map, visualize_token_wise_attention
+import matplotlib.patches as patches
+from matplotlib.colors import ListedColormap, BoundaryNorm
+from analysis.denoise_textual_file_wdet import process_single_sample_packed
+PATCH_SIZE = 32 #32x32 patches For Qwen3-VL
+import pickle
+use_act = True
+import math
+
+def plot_pseudo_label(image_path, layer2_attn, grounding_attn_o2i, grounding_attn, is_sink_token = None,
+                            save_path=None, figsize=(15, 4), fg_bbox = None):
+    """
+    Plot image and attention maps.
+
+    Args:
+        image_path: str, path to the image file
+        layer2_attn: ndarray (576,), layer 2 attention weights
+        grounding_attn_o2i: ndarray (576,), grounding attention object to image
+        grounding_attn: ndarray (576,), grounding attention
+        is_sink_token: ndarray (576,) in bool type to indicate whether this token is a sink token
+        save_path: str, optional path to save the plot
+        figsize: tuple, figure size
+    """
+    # Create figure with subplots
+    fig, axes = plt.subplots(1, 4, figsize=figsize, dpi=200)
+
+    # Plot 1: Original Image
+    if os.path.exists(image_path):
+        img = Image.open(image_path)
+        axes[0].imshow(img)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+    else:
+        axes[0].text(0.5, 0.5, f'Image not found:\n{image_path}',
+                     ha='center', va='center', transform=axes[0].transAxes)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+
+    # Reshape attention arrays to 24x24
+    if layer2_attn.ndim == 1:
+        square_size_hw = int(math.sqrt(len(layer2_attn.flatten())))
+        square_size = (square_size_hw, square_size_hw)
+    else:
+        square_size = layer2_attn.shape
+    layer2_attn_2d = layer2_attn.reshape(square_size)
+    grounding_attn_o2i_2d = grounding_attn_o2i.reshape(square_size)
+    grounding_attn_2d = grounding_attn.reshape(square_size)
+    if is_sink_token is not None:
+        is_sink_token_2d = is_sink_token.reshape(square_size)
+
+
+    # Plot 2: Special handling for Pseudo-Label map vs. continuous attention
+    # Check if the data is likely discrete pseudo-labels (-1, 0, 1)
+    if -1 in layer2_attn:
+        # --- Discrete Colormap for Pseudo-Labels ---
+        axes[1].set_title('Pseudo Label Map')
+        # Define the colors: Ignore (-1)->Gray, Background (0)->Dark Purple, Foreground (1)->Yellow
+        colors = ['gray', 'black', '#FFE699']
+        cmap_discrete = ListedColormap(colors)
+        # Define the boundaries for the colormap: [-1.5, -0.5), [-0.5, 0.5), [0.5, 1.5]
+        boundaries = [-1.5, -0.5, 0.5, 1.5]
+        norm = BoundaryNorm(boundaries, cmap_discrete.N)
+        im1 = axes[1].imshow(layer2_attn_2d, cmap=cmap_discrete, norm=norm, aspect='auto')
+        # Customize the colorbar for discrete values
+        cbar = plt.colorbar(im1, ax=axes[1], shrink=0.6)
+        cbar.set_ticks([-1, 0, 1])
+        cbar.set_ticklabels(['Ignore', 'BG', 'FG'])
+    else:
+        # --- Continuous Colormap for standard attention ---
+        axes[1].set_title('Layer 2 Attention')
+        im1 = axes[1].imshow(layer2_attn_2d, cmap='viridis', aspect='auto')
+        plt.colorbar(im1, ax=axes[1], shrink=0.6)
+
+    # Plot 3: Grounding Attention O2I
+    im2 = axes[2].imshow(grounding_attn_o2i_2d, cmap='inferno', aspect='auto')
+    axes[2].set_title('Grounding Attn O2I')
+    axes[2].axis('off')
+    plt.colorbar(im2, ax=axes[2], shrink=0.6)
+
+    # Plot 4: Grounding Attention Q2I
+    im3 = axes[3].imshow(grounding_attn_2d, cmap='inferno', aspect='auto')
+    axes[3].set_title('Grounding Attn')
+    axes[3].axis('off')
+    plt.colorbar(im3, ax=axes[3], shrink=0.6)
+
+    # Highlight sink tokens on the attention maps
+    if is_sink_token is not None:
+        # Find the coordinates of all sink tokens
+        sink_coords = np.argwhere(is_sink_token_2d)
+        for row, col in sink_coords:
+            # Define the properties of the highlight box
+            rect_args = {
+                'xy': (col - 0.5, row - 0.5), # Bottom-left corner of the cell
+                'width': 1,
+                'height': 1,
+                'linewidth': 2.5,
+                'edgecolor': '#7CFC00',  # A bright, "LawnGreen" color
+                'facecolor': 'none'
+            }
+            # Add a new rectangle patch to each of the three attention plots
+            #axes[1].add_patch(patches.Rectangle(**rect_args))
+            axes[2].add_patch(patches.Rectangle(**rect_args))
+            #axes[3].add_patch(patches.Rectangle(**rect_args))
+
+    if fg_bbox is not None:
+        fg_rect_arg = {
+            'xy': (fg_bbox[1] - 0.5, fg_bbox[0] - 0.5),  # Bottom-left corner of the cell
+            'width': fg_bbox[3] - fg_bbox[1] + 1,
+            'height': fg_bbox[2] - fg_bbox[0] + 1,
+            'linewidth': 2.5,
+            'edgecolor': '#FF4500',  # A bright, "OrangeRed" color
+            'facecolor': 'none'
+        }
+        axes[1].add_patch(patches.Rectangle(**fg_rect_arg))
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+
+    plt.show()
+    pass
+
+def replace_with_aura_layers(model, start_layer, num_layers):
+    """
+    Replace model.model.language_model.layers[start..start+num-1] in-place
+    with AuraDecoderLayer instances initialized from the original layers.
+
+    Since roi_mask=None is used during inference, forward output is
+    bit-for-bit identical to the original (Path A in AuraAttention).
+    """
+    layers = model.model.language_model.layers
+    total_layers = len(layers)
+    assert start_layer + num_layers <= total_layers, (
+        f"Cannot replace layers {start_layer}..{start_layer + num_layers - 1}: "
+        f"model only has {total_layers} layers"
+    )
+
+    print(f"Replacing layers {start_layer}..{start_layer + num_layers - 1} with AuraDecoderLayer")
+    for i in range(num_layers):
+        idx = start_layer + i
+        original_layer = layers[idx]
+        aura_layer = AuraDecoderLayer(model.config.text_config, layer_idx=idx)
+        missing = aura_layer.load_from_standard_layer(original_layer)
+        print(f"  AuraLayer[{i}] <- layers[{idx}]  (new params: {missing})")
+        # Move to same device/dtype as original
+        aura_layer = aura_layer.to(
+            device=next(original_layer.parameters()).device,
+            dtype=next(original_layer.parameters()).dtype,
+        )
+        layers[idx] = aura_layer
+    print("Layer replacement complete.")
+
+
+def eval_model_qwen(model, processor, questions, image_folder,
+                    ocr_cache_file=None, is_natural_image=False, save_path=None, debug_samples=10):
+
+    sample = 0
+
+    ocr_cache = {}
+    if ocr_cache_file and os.path.exists(ocr_cache_file):
+        print(f"Loading existing OCR cache from {ocr_cache_file}...")
+        try:
+            with open(ocr_cache_file, 'rb') as f:
+                ocr_cache = pickle.load(f)
+            print(f"Loaded {len(ocr_cache)} cached entries.")
+        except Exception as e:
+            print(f"Failed to load cache: {e}. Starting fresh.")
+
+
+    for line in tqdm(questions[sample:]):
+        if use_act: activations.clear()
+        image_file = line["image"]
+        qs = line["text"] if 'text' in line else line['prompt']
+
+        if is_natural_image:
+            qs = qs.replace('Answer the question using a single word or phrase.',
+                            'Output grounding bounding box related to the question in JSON.')
+        else:
+            if 'Answer the question using a single word or phrase.' not in qs:
+                qs += ' Answer the question using a single word or phrase.'
+
+        image = Image.open(os.path.join(image_folder, image_file)).convert('RGB')
+        processor.image_processor.max_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        processor.image_processor.min_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        image_size = image.size
+        image = expand2square(image, (127, 127, 127))
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": f"{qs}"
+                    },
+                    {
+                        "type": "image",
+                        "image": image
+                    }
+                ]
+            }
+        ]
+        # Preparation for inference
+        inputs = processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt"
+        )
+        inputs = inputs.to(model.device)
+
+        # Inference: Generation of the output
+        generated_ids = model.generate(**inputs, max_new_tokens=128, output_attentions=True)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+        print("Question:", qs)
+        print(output_text)
+
+        if use_act and activations.get('visual_mask', None) is None:
+            vision_start_token_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+            vision_end_token_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+            pos = inputs['input_ids'].tolist()[0].index(vision_start_token_id) + 1
+            pos_end = inputs['input_ids'].tolist()[0].index(vision_end_token_id)
+            activations['visual_mask'] = torch.zeros((1, inputs['input_ids'].shape[1]), dtype=torch.bool, device=device)
+            activations['visual_mask'][0, pos:pos_end] = True
+        output_shape = inputs['image_grid_thw'].cpu().numpy().squeeze(0)[1:] / 2
+        output_shape = np.array(output_shape, dtype=np.int32)
+        attn_weights_reduction = 'mean'  # 'none' or 'mean'
+        layers = model.config.text_config.num_hidden_layers
+        attn_weights_map = get_attn_weights_map(activations, layers=layers, reduction=attn_weights_reduction)
+
+
+        visualize_token_wise_attention(attn_weights_map, qk_type='output2images', attn_size=output_shape, layers=layers, rgb_image=image)
+
+        pass
+
+        if "Qwen3-VL-4B" in args.model_path and use_act:
+            if is_natural_image: ##TODO: find heads for natural image
+                required_heads = {20: [15, 21, 23], 21: [8, 10, 11]}
+                sink_heads = {20: [1, 24]}
+            else:
+                ##textual image
+                required_heads = {19: [24, 27], 24: [21, 23]}
+                sink_heads = {24: [6, 10]}
+            required_layers = list(required_heads.keys())
+        elif "Qwen3-VL-8B" in args.model_path and use_act:
+            if is_natural_image:  ##TODO: find heads for natural image
+                required_heads = {20: [21, 23], 21: [8, 10, 11]}
+                sink_heads = {21: [7]}
+            else:
+                ##textual image
+                required_heads = {19: [24, 27], 24: [21, 23]}
+                sink_heads = {24: [1, 2]}
+            required_layers = list(required_heads.keys())
+
+        grounding_attn_o2i = np.zeros(output_shape)
+        total_grounding_heads = sum([len(heads) for heads in required_heads.values()])
+        for layer, heads in required_heads.items():
+            attn_weights = attn_weights_map[f'layer{layer}']['output2images'].reshape(-1, output_shape[0], output_shape[1])
+            for head in heads:
+                grounding_attn_o2i += attn_weights[head]
+        grounding_attn_o2i = grounding_attn_o2i / total_grounding_heads
+
+        sink_attn = np.zeros(output_shape)
+        total_sink_heads = sum([len(heads) for heads in sink_heads.values()])
+        if total_sink_heads > 0:
+            for layer, heads in sink_heads.items():
+                attn_weights = attn_weights_map[f'layer{layer}']['output2images'].reshape(-1, output_shape[0], output_shape[1])
+                for head in heads:
+                    sink_attn += attn_weights[head]
+        sink_attn = sink_attn / total_sink_heads
+
+        ##denoise with ocr boxes
+        img_path = os.path.join(image_folder, image_file)
+        if not is_natural_image:
+            cached_boxes = ocr_cache.get(img_path, None)
+            item = {
+                'image': img_path,
+                'grounding_attn_o2i': grounding_attn_o2i,
+            }
+            new_item, _, _ = process_single_sample_packed((item, None, cached_boxes))
+            grounding_attn_o2i = new_item['grounding_attn_o2i']
+
+        pseudo_set_denoised = create_pseudo_labels(sink_attn, grounding_attn_o2i, sink_thresh=0.02, bg_coff=0.05, binary_coff=0.25,
+                                          original_image_size=image_size, for_vis=True, mask_known_bg = True)
+        plot_pseudo_label(
+            image_path=img_path,
+            layer2_attn=pseudo_set_denoised['labels'],
+            grounding_attn_o2i=grounding_attn_o2i,
+            grounding_attn=pseudo_set_denoised['grounding_attn'],
+            is_sink_token=pseudo_set_denoised['sink_token_mask'],
+            fg_bbox=pseudo_set_denoised['fg_bbox'],
+            save_path=f'{save_path}/{sample}_pseudo_label.png',
+        )
+        sample += 1
+        if sample >= debug_samples:
+            print('Debug samples limit reached. Stopping evaluation.')
+            break
+        pass
+
+import argparse
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, required=True, help="model name")
+    parser.add_argument("--model_path", type=str, required=True, help="model path")
+    parser.add_argument(
+        "--dataset", type=str, required=True,
+        choices=["ocr_vqa", "DocVQA", "textvqa", "gqa"],
+        help="ocr_vqa, DocVQA, textvqa, gqa"
+    )
+    parser.add_argument("--save_path", type=str, required=True, help="save directory")
+
+    parser.add_argument("--questions_path", type=str, required=True, help="path to .jsonl/.json questions file")
+    parser.add_argument("--images_path", type=str, required=True, help="path to images root folder, not include images/")
+    parser.add_argument("--ocr_cache_file", type=str, default=None, help="path to ocr cache .pkl (optional)")
+
+    parser.add_argument("--debug_samples", type=int, default=10, help="stop after N samples")
+    parser.add_argument("--aura_start_layer", type=int, default=24, help="first layer to replace with AuraDecoderLayer (twig_K)")
+    parser.add_argument("--aura_num_layers", type=int, default=3, help="number of layers to replace (twig_T)")
+    args = parser.parse_args()
+
+    device = "cuda"
+
+    model = Qwen3VLForConditionalGeneration.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+    ).eval().to(device)
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        padding_side="left",
+        use_fast=True,
+    )
+
+    # Step 1: Replace middle layers with AuraDecoderLayer (roi_mask=None → bit-for-bit identical)
+    replace_with_aura_layers(model, args.aura_start_layer, args.aura_num_layers)
+
+    # Step 2: Register hooks AFTER replacement so they attach to AuraAttention.self_attn
+    if use_act:
+        activations, hooks = register_hooks(model, is_qwen3vl=True)
+
+    is_natural_image = (args.dataset == "gqa")
+
+    if args.dataset == "textvqa":
+        args.images_path = args.images_path + "/textvqa/train_images"
+
+    print("args.images_path: ", args.images_path)
+    ocr_cache_file = os.path.expanduser(args.ocr_cache_file) if args.ocr_cache_file else None
+
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.questions_path), "r")]
+    print('Total questions loaded:', len(questions))
+    questions = [q for q in questions if args.dataset in os.path.join(args.images_path, q["image"])]
+    print('Questions after filtering by target dataset:', len(questions))
+
+
+    eval_model_qwen(
+        model=model,
+        processor=processor,
+        questions=questions,
+        image_folder=args.images_path,
+        ocr_cache_file=ocr_cache_file,
+        is_natural_image=is_natural_image,
+        save_path=args.save_path,
+        debug_samples=args.debug_samples,
+    )
+
+    print(f"dataset: {args.dataset}")
+    print(f"questions_path: {args.questions_path}")
+    print(f"images_path: {args.images_path}")
+    print(f"ocr_cache_file: {ocr_cache_file}")
+    print(f"is_natural_image: {is_natural_image}")
+    print(f"aura_start_layer: {args.aura_start_layer}")
+    print(f"aura_num_layers: {args.aura_num_layers}")

QWENVL-PRIVATE/analysis/analysis_auto.py

@@ -0,0 +1,588 @@
+import os
+import glob
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+
+# ==========================================
+# 1. DATA LOADING & PROCESSING
+# ==========================================
+
+def load_data_for_task(base_dir, task_name, model_pattern):
+    """
+    Loads pickle files matching the model_pattern for a specific task.
+    Returns:
+        - samples_data: dict of sample-level data (scores, visual tokens)
+        - metadata: dict containing 'e2e_latency', 'total_tokens', etc.
+    """
+    # Construct search pattern. logic handles wildcards if passed in model_pattern
+    search_pattern = os.path.join(base_dir, f"{task_name}_analysis_rank*_{model_pattern}*.pkl")
+    files = glob.glob(search_pattern)
+
+    if not files:
+        print(f"  [Warning] No files found for task '{task_name}' with pattern '*{model_pattern}*'")
+        return {}, {}
+
+    merged_samples = {}
+    aggregated_metadata = {"e2e_latency": 0, "total_tokens": 0}
+    found_metadata = False
+
+    for file_path in files:
+        try:
+            with open(file_path, "rb") as f:
+                data_chunk = pickle.load(f)
+
+                # Separate metric_dict from sample data
+                if 'metric_dict' in data_chunk:
+                    m_dict = data_chunk.pop('metric_dict')
+                    # Sum latencies/tokens if split across multiple files (rare but safe)
+                    aggregated_metadata['e2e_latency'] += m_dict.get('e2e_latency', 0)
+                    aggregated_metadata['total_tokens'] += m_dict.get('total_tokens', 0)
+                    found_metadata = True
+
+                merged_samples.update(data_chunk)
+        except Exception as e:
+            print(f"  [Error] Failed to load {file_path}: {e}")
+
+    if not found_metadata:
+        print(f"  [Info] No 'metric_dict' found in files for {task_name}.")
+
+    return merged_samples, aggregated_metadata
+
+
+def extract_metrics(task_data_map, task_meta_map):
+    """
+    Calculates average metrics (Throughput, Avg Visual Tokens) from raw data.
+    """
+    visual_tokens = []
+
+    # 1. Extract Visual Token Counts
+    for doc_id, content in task_data_map.items():
+        if "visual_token_num" in content:
+            val = content["visual_token_num"]
+            # Handle if it's a tensor/array or scalar
+            if hasattr(val, 'item'):
+                val = val.item()
+            visual_tokens.append(float(val))
+
+    avg_vis_tokens = np.mean(visual_tokens) if visual_tokens else 0
+
+    # 2. Calculate Throughput (Samples / Second)
+    # Total samples in this dataset
+    n_samples = len(task_data_map)
+
+    # Total time (e2e_latency is usually in seconds, ensure your logger units match)
+    total_time = task_meta_map.get('e2e_latency', 0)
+
+    if total_time > 0:
+        throughput = n_samples / total_time
+    else:
+        throughput = 0
+
+    return {
+        "avg_visual_tokens": avg_vis_tokens,
+        "throughput": throughput,
+        "n_samples": n_samples
+    }
+
+
+# ==========================================
+# 2. PLOTTING FUNCTIONS
+# ==========================================
+
+def plot_tradeoff_curves(results, save_dir):
+    """
+    Plots Accuracy vs. Efficiency (Throughput) and Accuracy vs. Visual Tokens.
+    Handles Baseline, Ours, and RoI curves with distinct styles.
+    """
+    labels = list(results.keys())
+
+    # Calculate Macro Averages (across all tasks for each config)
+    avg_accs = []
+    avg_throughputs = []
+    avg_tokens = []
+
+    for label in labels:
+        metrics = results[label]
+        # Calculate mean of values; ensure we handle potential empty lists safely
+        acc_values = [m['acc'] for m in metrics.values()]
+        th_values = [m['throughput'] for m in metrics.values()]
+        tok_values = [m['avg_visual_tokens'] for m in metrics.values()]
+
+        avg_accs.append(np.mean(acc_values) if acc_values else 0)
+        avg_throughputs.append(np.mean(th_values) if th_values else 0)
+        avg_tokens.append(np.mean(tok_values) if tok_values else 0)
+
+    # --- STYLE DEFINITIONS ---
+    # Customize colors and markers here
+    styles = {
+        'Baseline': {'marker': 'o', 'linestyle': '--', 'color': '#7f8c8d'},  # Grey, Dashed, Circle
+        'Ours': {'marker': '*', 'linestyle': '-', 'color': '#e74c3c'},  # Red, Solid, Star
+        'RoI': {'marker': 's', 'linestyle': '-', 'color': '#3498db'}  # Blue, Solid, Square
+    }
+
+    # Helper function to generate both plots without code duplication
+    def create_plot(x_data, x_label, file_name, title):
+        plt.figure(figsize=(10, 7))
+        sns.set_theme(style="whitegrid")  # Optional: makes grid look nice
+
+        # Identify unique prefixes (Baseline, Ours, RoI) from the config keys
+        prefixes = sorted(list(set([l.split('\n')[0] for l in labels])))
+
+        # Sort prefixes to ensure Legend order (e.g. Baseline -> Ours -> RoI) if desired
+        # You can manually force order: prefixes = ['Baseline', 'Ours', 'RoI']
+
+        for prefix in prefixes:
+            # Get style, default to black if unknown prefix
+            style = styles.get(prefix, {'marker': 'D', 'linestyle': '-.', 'color': 'black'})
+
+            # Filter data for this group
+            group_indices = [i for i, l in enumerate(labels) if l.startswith(prefix)]
+            if not group_indices: continue
+
+            g_x = [x_data[i] for i in group_indices]
+            g_y = [avg_accs[i] for i in group_indices]
+            g_labels = [labels[i] for i in group_indices]
+
+            # Sort points by X-axis (efficiency or tokens) to draw correct line
+            sorted_pairs = sorted(zip(g_x, g_y, g_labels), key=lambda x: x[0])
+            s_x, s_y, s_l = zip(*sorted_pairs)
+
+            # Plot Line
+            plt.plot(s_x, s_y,
+                     marker=style['marker'],
+                     linestyle=style['linestyle'],
+                     color=style['color'],
+                     label=prefix,
+                     linewidth=2,
+                     markersize=10 if style['marker'] != '*' else 14)  # Stars need to be larger
+
+            # Annotate Points (Labels like "Max-256")
+            for x, y, lbl in zip(s_x, s_y, s_l):
+                clean_lbl = lbl.split('\n')[-1]
+
+                # Dynamic offsets to prevent overlapping labels
+                offset = (0, 10)
+                if prefix == 'RoI':
+                    offset = (0, -15)
+                elif prefix == 'Ours':
+                    offset = (-10, 10)
+
+                plt.annotate(clean_lbl, (x, y), xytext=offset, textcoords='offset points',
+                             ha='center', fontsize=8, fontweight='bold',
+                             bbox=dict(boxstyle="round,pad=0.2", fc="white", alpha=0.6, ec="none"))
+
+        plt.title(title, fontsize=14)
+        plt.xlabel(x_label, fontweight='bold')
+        plt.ylabel('Average Accuracy (%)', fontweight='bold')
+        plt.grid(True, linestyle=':', alpha=0.6)
+        plt.legend(fontsize=10)
+        plt.tight_layout()
+
+        save_path = os.path.join(save_dir, file_name)
+        plt.savefig(save_path, dpi=300)
+        plt.close()
+        print(f"Saved plot: {save_path}")
+
+    # 1. Plot Efficiency
+    create_plot(avg_throughputs,
+                'Average Throughput (Samples / Second) \n← Slower | Faster →',
+                "acc_vs_efficiency.png",
+                "Accuracy vs. Efficiency Trade-off")
+
+    # 2. Plot Visual Tokens
+    create_plot(avg_tokens,
+                'Average Visual Tokens per Image \n← Cheaper | More Detailed →',
+                "acc_vs_visual_tokens.png",
+                "Accuracy vs. Visual Cost")
+
+# ==========================================
+# 3. EXISTING UTILS (Preserved)
+# ==========================================
+def print_threshold_ratios(task_data_map, threshold):
+    print(f"\n{'=' * 20} Threshold Analysis (Score > {threshold}) {'=' * 20}")
+    print(f"{'Dataset Name':<20} | {'Total Samples':<15} | {'Above Thr':<10} | {'Ratio':<10}")
+    print("-" * 65)
+    for task_name, content in task_data_map.items():
+        scores = content['scores']
+        if len(scores) == 0: continue
+        above_mask = scores > threshold
+        count_above = np.sum(above_mask)
+        ratio = count_above / len(scores)
+        print(f"{task_name:<20} | {len(scores):<15} | {count_above:<10} | {ratio:.2%}")
+    print("-" * 65 + "\n")
+
+
+def extract_scores_list(data_dict):
+    scores = []
+    for doc_id, content in data_dict.items():
+        if "high_res_pred_score" in content:
+            val = content["high_res_pred_score"]
+            if isinstance(val, np.ndarray): val = val.flatten()[0]
+            scores.append(float(val))
+    return np.array(scores)
+
+
+def plot_cost_benefit_combo(results, task_name, save_dir):
+    """
+    Better Idea: Combines Relative Efficiency (Bars) with Accuracy Gain (Points/Lines)
+    on a dual-axis chart.
+    """
+    # 1. Parse Data (Similar to above, but also gets Accuracy)
+    data_map = {}
+    for label, metrics in results.items():
+        method, resolution = label.split('\n')
+        if resolution not in data_map: data_map[resolution] = {}
+
+        task_data = metrics.get(task_name, {'throughput': 0, 'acc': 0})
+        data_map[resolution][method] = {
+            'speed': task_data['throughput'],
+            'acc': task_data['acc']
+        }
+
+    resolutions = sorted(data_map.keys(), key=lambda x: int(x.split('-')[-1]))
+    methods = ['Baseline', 'Ours', 'RoI']
+    colors = {'Baseline': '#95a5a6', 'Ours': '#e74c3c', 'RoI': '#3498db'}
+
+    x = np.arange(len(resolutions))
+    width = 0.25
+
+    fig, ax1 = plt.subplots(figsize=(11, 7))
+    ax2 = ax1.twinx()  # Create secondary y-axis for Accuracy Gain
+
+    # Loop for Bars (Efficiency) AND Points (Accuracy)
+    for i, method in enumerate(methods):
+        eff_ratios = []
+        acc_gains = []
+
+        for res in resolutions:
+            base_speed = data_map[res].get('Baseline', {}).get('speed', 1.0)
+            base_acc = data_map[res].get('Baseline', {}).get('acc', 0.0)
+
+            curr = data_map[res].get(method, {'speed': 0, 'acc': 0})
+
+            # Efficiency Ratio
+            ratio = (curr['speed'] / base_speed) * 100 if base_speed > 0 else 0
+            eff_ratios.append(ratio)
+
+            # Accuracy Delta (Gain over Baseline of SAME resolution)
+            delta = curr['acc'] - base_acc
+            acc_gains.append(delta)
+
+        offset = (i - 1) * width
+
+        # PLOT 1: Efficiency Bars (Left Axis)
+        ax1.bar(x + offset, eff_ratios, width, color=colors[method], alpha=0.3, label=f"{method} (Speed)")
+
+        # PLOT 2: Accuracy Gain Markers (Right Axis)
+        # We plot these slightly offset to align with bars, or connected with lines
+        if method == 'Baseline':
+            # Plot baseline as flat line at 0
+            ax2.plot(x + offset, acc_gains, color=colors[method], marker='_', markersize=20, linestyle='None')
+        else:
+            # Plot gains as distinct markers
+            ax2.plot(x + offset, acc_gains, color=colors[method], marker='D', markersize=8,
+                     linestyle='-', linewidth=2, label=f"{method} (Acc Gain)")
+
+            # Annotate specific gains
+            for j, gain in enumerate(acc_gains):
+                ax2.annotate(f"+{gain:.1f}", (x[j] + offset, gain), xytext=(0, 8),
+                             textcoords='offset points', ha='center', color=colors[method], fontweight='bold')
+
+    # Formatting
+    ax1.set_ylabel("Relative Efficiency (%)", fontsize=12, fontweight='bold')
+    ax1.set_ylim(0, 120)
+    ax1.set_xlabel("Resolution", fontsize=12)
+    ax1.set_xticks(x)
+    ax1.set_xticklabels(resolutions)
+
+    ax2.set_ylabel("Accuracy Gain over Baseline (Points)", fontsize=12, fontweight='bold')
+    # Set symmetric limits or auto limits for acc gain
+    ax2.axhline(0, color='gray', linestyle=':', linewidth=1)
+
+    plt.title(f"Cost-Benefit Analysis: {task_name}\n(Bars = Speed %, Points = Acc Gain)", fontsize=14)
+
+    # Custom Legend
+    lines, labels = ax1.get_legend_handles_labels()
+    lines2, labels2 = ax2.get_legend_handles_labels()
+    ax1.legend(lines + lines2, labels + labels2, loc='upper left')
+
+    plt.tight_layout()
+    save_path = os.path.join(save_dir, f"{task_name}_cost_benefit.png")
+    plt.savefig(save_path, dpi=300)
+    plt.close()
+    print(f"Saved Cost-Benefit plot: {save_path}")
+
+# ==========================================
+# 4. MAIN EXECUTION
+# ==========================================
+
+if __name__ == "__main__":
+    BASE_DIR = "../archive_data/"
+    SAVE_DIR = "./analysis_plots/"
+    os.makedirs(SAVE_DIR, exist_ok=True)
+
+    # --- CONFIGURATION: Define Models and Manual Accuracies ---
+    # Keys: Display Label (First line is Group, second is variant)
+    # Values: 'pattern': filename part to match, 'accs': dict of {task: accuracy}
+    TARGET_TASKS = [
+        "docvqa_val",
+        "infovqa_val",
+        "textvqa_val",
+        "chartqa",  # Added from table
+        "ocrbench"  # Added from table
+        # "vstar_bench",
+        # "mmerealworld_lite",
+    ]
+
+    # --- CONFIGURATION: Define Models and Manual Accuracies ---
+    # Patterns are derived from: Visual_Tokens * 28 * 28
+    # 128 tokens = 100352 pixels
+    # 256 tokens = 200704 pixels
+    # 576 tokens = 451584 pixels
+    # 1024 tokens = 802816 pixels
+    # 2048 tokens = 1605632 pixels
+
+    MODEL_CONFIGS = {
+        # --- BASELINE CONFIGURATIONS ---
+        # 'Baseline\nMax-256': {
+        #     'pattern': 'Qwen2.5-VL-7B-Instruct_min100352_max200704',
+        #     'accs': {
+        #         'docvqa_val': 81.2,
+        #         'chartqa': 76.8,
+        #         'ocrbench': 77.8,
+        #         'infovqa_val': 54.0,
+        #         'textvqa_val': 75.5,
+        #         'vstar_bench': 53.4,
+        #         'mmerealworld_lite': 27.6
+        #     }
+        # },
+        'Baseline\nMax-576': {
+            'pattern': 'Qwen2.5-VL-7B-Instruct_min200704_max451584',
+            'accs': {
+                'docvqa_val': 92.0,
+                'chartqa': 83.0,
+                'ocrbench': 82.8,
+                'infovqa_val': 70.1,
+                'textvqa_val': 81.1,
+                'vstar_bench': 64.4,
+                'mmerealworld_lite': 35.4
+            }
+        },
+        'Baseline\nMax-1024': {
+            'pattern': 'Qwen2.5-VL-7B-Instruct_min200704_max802816',
+            'accs': {
+                'docvqa_val': 93.9,
+                'chartqa': 84.6,
+                'ocrbench': 84.4,
+                'infovqa_val': 76.5,
+                'textvqa_val': 82.8,
+                'vstar_bench': 70.2,
+                'mmerealworld_lite': 37.5
+            }
+        },
+        'Baseline\nMax-2048': {
+            'pattern': 'Qwen2.5-VL-7B-Instruct_min200704_max1605632',
+            'accs': {
+                'docvqa_val': 94.7,
+                'chartqa': 84.5,
+                'ocrbench': 84.8,
+                'infovqa_val': 80.4,
+                'textvqa_val': 82.9,
+                'vstar_bench': 72.3,
+                'mmerealworld_lite': 39.3
+            }
+        },
+        #
+        # # # # --- YOUR "OURS" CONFIGURATIONS (Update these patterns as needed) ---
+        # 'Ours\nMax-256': {
+        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min100352_max200704',
+        #     'accs': {
+        #         'docvqa_val': 89.1,
+        #         'chartqa': 82.5,
+        #         'ocrbench': 83.8,
+        #         'infovqa_val': 69.6,
+        #         'textvqa_val': 81.8,
+        #         'vstar_bench': 67.5,
+        #         'mmerealworld_lite': 37.3
+        #     }
+        # },
+        # 'Ours\nMax-576': {
+        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min200704_max451584',
+        #     'accs': {
+        #         'docvqa_val': 93.9,
+        #         'chartqa': 84.2,
+        #         'ocrbench': 84.8,
+        #         'infovqa_val': 78.5,
+        #         'textvqa_val': 82.8,
+        #         'vstar_bench': 80.1,
+        #         'mmerealworld_lite': 43.4
+        #     }
+        # },
+        # 'Ours\nMax-1024': {
+        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min200704_max802816',
+        #     'accs': {
+        #         'docvqa_val': 94.6,
+        #         'chartqa': 85.4,
+        #         'ocrbench': 84.8,
+        #         'infovqa_val': 81.5,
+        #         'textvqa_val': 83.7,
+        #         'vstar_bench': 83.3,
+        #         'mmerealworld_lite': 46.5
+        #     }
+        # },
+        # 'Ours\nMax-1536': {
+        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled_min200704_max1204224',
+        #     'accs': {
+        #         'docvqa_val': 94.8,
+        #         'chartqa': 85.4,
+        #         'ocrbench': 85.4,
+        #         'infovqa_val': 82.4,
+        #         'textvqa_val': 83.6,
+        #         'vstar_bench': 88.0,
+        #         'mmerealworld_lite': 47.2
+        #     }
+        # },
+        ###--- "OURS-RoI" CONFIGURATIONS ---
+        'RoI\nMax-256': {
+            'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min100352_max200704',
+            'accs': {
+                'docvqa_val': 89.2, 'chartqa': 83.0, 'ocrbench': 83.6, 'infovqa_val': 69.9,
+                'textvqa_val': 81.4, 'vstar_bench': 67.5, 'mmerealworld_lite': 37.6
+            }
+        },
+        'RoI\nMax-576': {
+            'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min200704_max451584',
+            'accs': {
+                'docvqa_val': 94.1, 'chartqa': 85.8, 'ocrbench': 84.9, 'infovqa_val': 79.6,
+                'textvqa_val': 82.9, 'vstar_bench': 80.6, 'mmerealworld_lite': 44.6
+            }
+        },
+        'RoI\nMax-1024': {
+            'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min200704_max802816',
+            'accs': {
+                'docvqa_val': 94.7, 'chartqa': 86.3, 'ocrbench': 84.8, 'infovqa_val': 82.0,
+                'textvqa_val': 83.9, 'vstar_bench': 82.2, 'mmerealworld_lite': 46.7
+            }
+        },
+        # 'RoI\nMax-1536': {
+        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled_min200704_max1204224',
+        #     'accs': {
+        #         'docvqa_val': 94.9, 'chartqa': 85.8, 'ocrbench': 84.8, 'infovqa_val': 83.2,
+        #         'textvqa_val': 83.9, 'vstar_bench': 88.0, 'mmerealworld_lite': 47.0
+        #     }
+        # },
+
+        # # # --- YOUR "OURS" CONFIGURATIONS (Update these patterns as needed) ---
+        'Ours\nMax-1024': {
+            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min200704_max802816',
+            'accs': {
+                'docvqa_val': 94.8,
+                'chartqa': 86.4,
+                'ocrbench': 85.4,
+                'infovqa_val': 81.2,
+                'textvqa_val': 84.1,
+                'vstar_bench': 83.8,
+                'mmerealworld_lite': 47.3
+            }
+        },
+        'Ours\nMax-576': {
+            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min200704_max451584',
+            'accs': {
+                'docvqa_val': 94.0,
+                'chartqa': 85.4,
+                'ocrbench': 85.4,
+                'infovqa_val': 79.4,
+                'textvqa_val': 83.5,
+                'vstar_bench': 79.6,
+                'mmerealworld_lite': 45.7
+            }
+        },
+        'Ours\nMax-384': {
+            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min100352_max301056',
+            'accs': {
+                'docvqa_val': 93.4,
+                'chartqa': 84.2,
+                'ocrbench': 85.3,
+                'infovqa_val': 75.7,
+                'textvqa_val': 82.7,
+                'vstar_bench': 73.8,
+                'mmerealworld_lite': 41.2
+            }
+        },
+        'Ours\nMax-256': {
+            'pattern': 'qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-stage3_min100352_max200704',
+            'accs': {
+                'docvqa_val': 90.4,
+                'chartqa': 82.7,
+                'ocrbench': 83.5,
+                'infovqa_val': 71.7,
+                'textvqa_val': 81.8,
+                'vstar_bench': 68.6,
+                'mmerealworld_lite': 37.9
+            }
+        },
+        # 'Ours\nMax-1536': {
+        #     'pattern': 'qwen2_5vl-7b-roi-K18T3-post7k-repos-full2_min200704_max1204224',
+        #     'accs': {
+        #         'docvqa_val': 95.0,
+        #         'chartqa': 86.4,
+        #         'ocrbench': 84.8,
+        #         'infovqa_val': 82.6,
+        #         'textvqa_val': 84.3,
+        #         'vstar_bench': 86.9,
+        #         'mmerealworld_lite': 48.2
+        #     }
+        # },
+
+    }
+    print("--- Loading Data for Comparison ---")
+
+    # Store all results here: { "Baseline\nMax-256": { "textvqa": {throughput, tokens, acc}, ... } }
+    overall_results = {}
+
+    for label, config in MODEL_CONFIGS.items():
+        print(f"Processing config: {label.replace(chr(10), ' ')}")  # chr(10) is newline
+        overall_results[label] = {}
+
+        # Load data for each task for this specific config
+        for task in TARGET_TASKS:
+            # 1. Load Raw Data
+            raw_data, metadata = load_data_for_task(BASE_DIR, task, config['pattern'])
+
+            if not raw_data:
+                # Fill dummy if missing to avoid crash, or skip
+                overall_results[label][task] = {'throughput': 0, 'avg_visual_tokens': 0, 'acc': 0}
+                continue
+
+            # 2. Extract Metrics (Speed, Tokens)
+            metrics = extract_metrics(raw_data, metadata)
+
+            # 3. Add Manual Accuracy
+            acc = config['accs'].get(task, 0.0)  # Default 0 if task missing in config
+
+            # 4. Store Scores for Histograms (Optional, if you want detailed score plots per config)
+            # scores = extract_scores_list(raw_data)
+
+            # Save combined
+            overall_results[label][task] = {
+                'throughput': metrics['throughput'],
+                'avg_visual_tokens': metrics['avg_visual_tokens'],
+                'acc': acc,
+                # 'scores': scores # Uncomment if you want to pass this to the histogram plotter
+            }
+
+            print(
+                f"  > {task}: {metrics['throughput']:.1f} samp/s, {metrics['avg_visual_tokens']:.0f} tokens, Acc: {acc}")
+
+    # --- PLOT COMPARISONS ---
+    if overall_results:
+        print("\n--- Generating Comparison Plots ---")
+        plot_tradeoff_curves(overall_results, SAVE_DIR)
+
+    print("\nDone! Plots saved to:", SAVE_DIR)
+
+    for i in range(len(TARGET_TASKS)):
+        plot_cost_benefit_combo(overall_results, TARGET_TASKS[i], SAVE_DIR)

QWENVL-PRIVATE/analysis/analysis_high_res_prediction.py

@@ -0,0 +1,190 @@
+import os
+import glob
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+
+def print_threshold_ratios(task_data_map, threshold):
+    """
+    Calculates and prints the ratio of samples above the given threshold.
+    """
+    print(f"\n{'=' * 20} Threshold Analysis (Score > {threshold}) {'=' * 20}")
+    print(f"{'Dataset Name':<20} | {'Total Samples':<15} | {'Above Thr':<10} | {'Ratio':<10}")
+    print("-" * 65)
+
+    for task_name, scores in task_data_map.items():
+        if len(scores) == 0:
+            continue
+
+        # Boolean mask: True if score > threshold
+        above_mask = scores > threshold
+        count_above = np.sum(above_mask)
+        total_samples = len(scores)
+        ratio = count_above / total_samples
+
+        # Print formatted row
+        print(f"{task_name:<20} | {total_samples:<15} | {count_above:<10} | {ratio:.2%}")
+    print("-" * 65 + "\n")
+
+def load_data_for_task(base_dir, task_name, min_pixels=200704, max_pixels=451584, model_name="qwen2_5vl-3b-roi-K24T3-stage3"):
+    """
+    Loads all pickle files for a specific task (merging different ranks).
+    """
+    search_pattern = os.path.join(base_dir, f"{task_name}_analysis_rank*_{model_name}_min{min_pixels}_max{max_pixels}.pkl")
+    print("Searching for files with pattern:", search_pattern)
+    files = glob.glob(search_pattern)
+
+    if not files:
+        print(f"Warning: No files found for task '{task_name}' in {base_dir}")
+        return {}
+
+    merged_data = {}
+    print(f"Loading {len(files)} files for task: {task_name}...")
+
+    for file_path in files:
+        try:
+            with open(file_path, "rb") as f:
+                data_chunk = pickle.load(f)
+                merged_data.update(data_chunk)
+        except Exception as e:
+            print(f"Error loading {file_path}: {e}")
+
+    return merged_data
+
+
+def extract_scores(data_dict):
+    """
+    Extracts the 'high_res_pred_score' from the dictionary into a flat list.
+    """
+    scores = []
+    for doc_id, content in data_dict.items():
+        if "high_res_pred_score" in content:
+            val = content["high_res_pred_score"]
+            if isinstance(val, np.ndarray):
+                val = val.flatten()[0]
+            scores.append(float(val))
+    return np.array(scores)
+
+
+def plot_single_task(task_name, scores, save_dir):
+    """
+    Plots the distribution for a SINGLE task and saves it to a separate file.
+    """
+    if len(scores) == 0:
+        return
+
+    plt.figure(figsize=(8, 6))
+    sns.set_theme(style="whitegrid")
+
+    # Plot specific color for individual plots
+    color = sns.color_palette("viridis", 1)[0]
+
+    sns.histplot(
+        scores,
+        kde=True,
+        stat="probability",  # Normalized ratio
+        element="step",
+        bins=50,
+        color=color,
+        label=f"{task_name} (n={len(scores)})"
+    )
+
+    # Calculate stats for title
+    mean_val = np.mean(scores)
+    std_val = np.std(scores)
+
+    plt.title(f"Distribution: {task_name}\n(Mean={mean_val:.2f}, Std={std_val:.2f})")
+    plt.xlabel("Prediction Score")
+    plt.ylabel("Density (Normalized)")
+    plt.legend()
+    plt.tight_layout()
+    plt.show()
+
+    save_path = os.path.join(save_dir, f"{task_name}_dist.png")
+    plt.savefig(save_path, dpi=300)
+    plt.close()  # Close figure to free memory
+    print(f"Saved individual plot: {save_path}")
+
+
+def plot_overall_comparison(task_data_map, save_path):
+    """
+    Plots ALL tasks on the SAME figure for comparison.
+    """
+    plt.figure(figsize=(12, 7))
+    sns.set_theme(style="whitegrid")
+
+    # Use a distinct palette for multiple lines
+    palette = sns.color_palette("bright", len(task_data_map))
+
+    for i, (task_name, scores) in enumerate(task_data_map.items()):
+        if len(scores) == 0:
+            continue
+
+        sns.histplot(
+            scores,
+            label=f"{task_name}",
+            kde=True,
+            element="step",
+            stat="probability",
+            common_norm=False,
+            bins=50,
+            alpha=0.3,
+            color=palette[i]
+        )
+
+    plt.title("Overall Comparison: High Res Prediction Scores")
+    plt.xlabel("Prediction Score")
+    plt.ylabel("Ratio Percentage")
+    plt.legend(title="Datasets")
+    plt.tight_layout()
+    plt.show()
+    plt.savefig(save_path, dpi=300)
+    plt.close()
+    print(f"Saved overall comparison plot: {save_path}")
+
+
+if __name__ == "__main__":
+    # 1. SETUP
+    BASE_DIR = "./excluded/"
+    # Directory to save the images
+    SAVE_DIR = "./excluded/analysis_plots/"
+    THRESHOLD = 0.1
+
+    # Ensure save directory exists
+    os.makedirs(SAVE_DIR, exist_ok=True)
+
+    TARGET_TASKS = ["textvqa_val", "infovqa_val", "chartqa", "ocrbench",
+                    "docvqa_val","vstar_bench", "hrbench8k","mmerealworld_lite"]
+    # 2. LOAD DATA
+    task_scores_map = {}
+
+    print("--- Loading Data ---")
+    for task in TARGET_TASKS:
+        raw_dict = load_data_for_task(BASE_DIR, task)
+        if raw_dict:
+            scores = extract_scores(raw_dict)
+            task_scores_map[task] = scores
+
+    if not task_scores_map:
+        print("No data found. Exiting.")
+        exit()
+
+    # 3. PRINT RATIO STATISTICS (NEW FEATURE)
+    print_threshold_ratios(task_scores_map, THRESHOLD)
+    print_threshold_ratios(task_scores_map, 0.05)
+    print_threshold_ratios(task_scores_map, 0.025)
+
+
+    print("\n--- Generating Plots ---")
+
+    # 3. PLOT INDIVIDUAL (One file per task)
+    for task_name, scores in task_scores_map.items():
+        plot_single_task(task_name, scores, SAVE_DIR)
+
+    # # 4. PLOT OVERALL (One file with all tasks)
+    # overall_path = os.path.join(SAVE_DIR, "overall_comparison.png")
+    # plot_overall_comparison(task_scores_map, overall_path)
+    #
+    # print("\nDone! All plots saved to:", SAVE_DIR)

QWENVL-PRIVATE/analysis/check_gate_label.py

@@ -0,0 +1,195 @@
+import pickle
+import matplotlib.pyplot as plt
+import os
+import math
+from PIL import Image
+import textwrap
+import string
+from collections import Counter
+import random
+
+def _normalize_answer(text):
+    """Helper to normalize text for comparison (lower case, strip punctuation)."""
+    if not isinstance(text, str): return str(text)
+    text = text.lower().strip()
+    # Remove punctuation
+    text = text.translate(str.maketrans('', '', string.punctuation))
+    return text
+
+
+def plot_single_sample(item, image_root_folder, save_path=None, show_plot=True):
+    """
+    Plots a single sample's image and resolution consistency chart.
+
+    Args:
+        item: A dictionary containing 'image', 'res', 'ms_answers', 'text' (GT), etc.
+        image_root_folder: Path to the root folder where images are stored.
+        save_path: If provided, saves the figure to this path.
+        show_plot: If True, calls plt.show().
+    """
+    high_res_group_num = 2
+    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
+
+    # --- Left Side: The Image ---
+    img_filename = item['image']
+
+    # Handle cases where image path might be absolute or relative
+    if os.path.isabs(img_filename):
+        img_path = img_filename
+    else:
+        img_path = os.path.join(image_root_folder, img_filename)
+
+    try:
+        img = Image.open(img_path).convert('RGB')
+        axes[0].imshow(img)
+        axes[0].axis('off')
+        axes[0].set_title(f"ID: {item['id']}\nGT: {item['text']}", fontsize=10)
+    except Exception as e:
+        axes[0].text(0.5, 0.5, f"Image Not Found:\n{img_filename}",
+                     ha='center', va='center', fontsize=12)
+        axes[0].axis('off')
+
+    # --- Right Side: Resolution vs Answer Flow ---
+    resolutions = item['res']  # e.g., [64, 128, 256...]
+    answers = item['ms_answers']  # e.g., ["dog", "dog", "cat"...]
+    # Normalize all answers for fair comparison
+    norm_answers = [_normalize_answer(a) for a in answers]
+    if len(norm_answers) >= high_res_group_num:
+        high_res_group = norm_answers[-high_res_group_num:]
+    else:
+        high_res_group = norm_answers  # Fallback for short lists
+    # Find the consensus answer
+    counts = Counter(high_res_group)
+    target_answer, freq = counts.most_common(1)[0]
+
+    # --- RULE 3: UNSTABLE / HOPELESS ---
+    # If the high-res views strongly disagree (e.g., ["A", "B", "C"]), ignore sample.
+    # Logic: If frequency is 1 (all different), it's unstable.
+    if len(high_res_group) >= high_res_group_num and freq < 2:
+        print('Skipping unstable sample ID:', item['id'])
+        plt.close(fig)
+        return # Skip plotting
+
+    # --- STEP 2: Categorize Indices ---
+    # Which resolutions got it "Right" vs "Wrong"?
+    correct_indices = [i for i, ans in enumerate(norm_answers) if (target_answer in ans or ans in target_answer)]
+    wrong_indices = [i for i, ans in enumerate(norm_answers) if (target_answer not in ans and ans not in target_answer)]
+
+    # --- RULE 1: EASY / CONSISTENT CASE ---
+    # If EVERYTHING is correct (or almost everything, specifically low res),
+    # we want to teach efficiency.
+    if len(wrong_indices) == 0:
+        print('Skipping consistent sample ID:', item['id'])
+        plt.close(fig)
+        return
+
+    # Clean answers for display (wrap long text)
+    wrapped_answers = ['\n'.join(textwrap.wrap(str(a), width=20)) for a in answers]
+
+    # Assign colors: Green if matches Ground Truth, Red if not (if GT exists), else Blue
+    gt = str(item['text']).lower().strip() if item['text'] else None
+
+    # Simple heuristic to determine "Final Consensus" if GT is missing or complex
+    final_answer = str(answers[-1]).lower().strip()
+
+    if target_answer not in gt and gt not in target_answer:
+        print('Skipping sample with complex/no GT ID:', item['id'])
+        plt.close(fig)
+        return
+
+    bar_colors = []
+    for ans in answers:
+        ans_clean = str(ans).lower().strip()
+        # Your specific containment logic
+        if gt and (gt in ans_clean or ans_clean in gt):
+            bar_colors.append('#2ca02c')  # Green (Correct)
+        elif (final_answer in ans_clean or ans_clean in final_answer):
+            bar_colors.append('#1f77b4')  # Blue (Consistent with final)
+        else:
+            bar_colors.append('#d62728')  # Red (Inconsistent/Wrong)
+
+    # Create a step plot or scatter plot
+    y_pos = range(len(resolutions))
+
+    # Draw connecting lines first
+    axes[1].plot(y_pos, [0] * len(y_pos), color='gray', alpha=0.3, zorder=1)  # Baseline
+
+    # Draw scattered points with labels
+    axes[1].scatter(y_pos, [0] * len(y_pos), s=100, c=bar_colors, zorder=2)
+
+    # Annotate texts
+    for idx, (res, ans_text, color) in enumerate(zip(resolutions, wrapped_answers, bar_colors)):
+        # Alternate height to avoid overlapping text
+        vertical_offset = 0.2 if idx % 2 == 0 else -0.2
+
+        # Highlight changed answers
+        font_weight = 'bold' if idx > 0 and answers[idx] != answers[idx - 1] else 'normal'
+
+        axes[1].text(idx, vertical_offset, ans_text,
+                     ha='center',
+                     va='center' if vertical_offset == 0 else ('bottom' if vertical_offset > 0 else 'top'),
+                     fontsize=9, fontweight=font_weight, color='black',
+                     bbox=dict(boxstyle="round,pad=0.3", fc="white", ec=color, alpha=0.9))
+
+        axes[1].text(idx, -0.05, f"{res}px", ha='center', va='top', fontsize=8, color='gray')
+
+    # Formatting right ax
+    axes[1].set_title(f"Prompt: {item['prompt'][:60]}...", fontsize=10)
+    axes[1].set_xlim(-0.5, len(resolutions) - 0.5)
+    axes[1].set_ylim(-1, 1)
+    axes[1].axis('off')  # Hide standard axes
+
+    # Add legend/info
+    info_text = "Logic:\nGreen = Matches GT\nBlue = Matches Final Res\nRed = Mismatch"
+    axes[1].text(0.95, 0.05, info_text, transform=axes[1].transAxes,
+                 fontsize=8, ha='right', va='bottom', bbox=dict(facecolor='white', alpha=0.5))
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path)
+
+    if show_plot:
+        plt.show()
+
+    plt.close(fig)
+
+
+def visualize_resolution_consistency(pkl_path, image_root_folder, save_dir='../excluded/vis_results', num_samples=10):
+    """
+    Visualizes samples from the pickle file to check consistency hypothesis.
+    """
+
+    # 1. Load Data
+    print(f"Loading data from {pkl_path}...")
+    try:
+        with open(pkl_path, 'rb') as f:
+            data = pickle.load(f)
+    except Exception as e:
+        print(f"Error loading pickle: {e}")
+        return
+
+    # Convert dict to list for easier slicing
+    # Note: Retained your specific slicing logic
+    samples = list(data.values())[150:]
+    selected_samples = samples
+
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    print(f"Generating plots for {len(selected_samples)} samples...")
+
+    # 3. Plotting Loop
+    for i, item in enumerate(selected_samples):
+        plot_single_sample(item, image_root_folder, save_path=None, show_plot=True)
+
+    print(f"Done! Check the '{save_dir}' folder.")
+
+# --- Usage ---
+if __name__ == "__main__":
+    # Update these paths to match your setup
+    PKL_FILE = "/home/yuheng/datasets/roi_vcot68k_ms.pkl"
+    IMAGE_ROOT = "/home/yuheng/datasets"
+
+    # Run visualization
+    visualize_resolution_consistency(PKL_FILE, IMAGE_ROOT, num_samples=50)

QWENVL-PRIVATE/analysis/check_pseudo_label.py

@@ -0,0 +1,630 @@
+import os
+import json
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.colors import ListedColormap, BoundaryNorm
+import torch
+import torchvision.transforms.functional as torchvision_F
+from tqdm import tqdm
+#from denoise_textual_file import refine_pseudo_labels_with_doctr
+# Add the parent directory to sys.path to import qwen_src
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from qwen_src.mm_utils import expand2square, create_pseudo_labels, get_foreground_mask
+
+def plot_attention_analysis(image_path, sink_attn, grounding_attn_o2i, grounding_attn, is_sink_token = None,
+                            save_path=None, figsize=(15, 4), fg_bbox = None):
+    """
+    Plot image and attention maps.
+
+    Args:
+        image_path: str, path to the image file
+        sink_attn: ndarray (576,), layer 2 attention weights
+        grounding_attn_o2i: ndarray (576,), grounding attention object to image
+        grounding_attn: ndarray (576,), grounding attention
+        is_sink_token: ndarray (576,) in bool type to indicate whether this token is a sink token
+        save_path: str, optional path to save the plot
+        figsize: tuple, figure size
+    """
+    # Create figure with subplots
+    fig, axes = plt.subplots(1, 4, figsize=figsize, dpi=200)
+    h,w = sink_attn.shape[0], sink_attn.shape[1] if len(sink_attn.shape) > 1 else 24
+    # Plot 1: Original Image
+    if os.path.exists(image_path):
+        img = Image.open(image_path)
+        axes[0].imshow(img)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+    else:
+        axes[0].text(0.5, 0.5, f'Image not found:\n{image_path}',
+                     ha='center', va='center', transform=axes[0].transAxes)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+
+    # Reshape attention arrays to 24x24
+    sink_attn_2d = sink_attn.reshape(h, w)
+    grounding_attn_o2i_2d = grounding_attn_o2i.reshape(h, w)
+    grounding_attn_2d = grounding_attn.reshape(h, w)
+    if is_sink_token is not None:
+        is_sink_token_2d = is_sink_token.reshape(h, w)
+
+
+    # Plot 2: Special handling for Pseudo-Label map vs. continuous attention
+    # Check if the data is likely discrete pseudo-labels (-1, 0, 1)
+    if -1 in sink_attn:
+        # --- Discrete Colormap for Pseudo-Labels ---
+        axes[1].set_title('Pseudo Label Map')
+        # Define the colors: Ignore (-1)->Gray, Background (0)->Dark Purple, Foreground (1)->Yellow
+        colors = ['gray', 'black', '#FFE699']
+        cmap_discrete = ListedColormap(colors)
+        # Define the boundaries for the colormap: [-1.5, -0.5), [-0.5, 0.5), [0.5, 1.5]
+        boundaries = [-1.5, -0.5, 0.5, 1.5]
+        norm = BoundaryNorm(boundaries, cmap_discrete.N)
+        im1 = axes[1].imshow(sink_attn_2d, cmap=cmap_discrete, norm=norm, aspect='auto')
+        # Customize the colorbar for discrete values
+        cbar = plt.colorbar(im1, ax=axes[1], shrink=0.6)
+        cbar.set_ticks([-1, 0, 1])
+        cbar.set_ticklabels(['Ignore', 'BG', 'FG'])
+    else:
+        # --- Continuous Colormap for standard attention ---
+        axes[1].set_title('Layer 2 Attention')
+        im1 = axes[1].imshow(sink_attn_2d, cmap='viridis', aspect='auto')
+        plt.colorbar(im1, ax=axes[1], shrink=0.6)
+
+    # Plot 3: Grounding Attention O2I
+    im2 = axes[2].imshow(grounding_attn_o2i_2d, cmap='inferno', aspect='auto')
+    axes[2].set_title('Grounding Attn O2I')
+    axes[2].axis('off')
+    plt.colorbar(im2, ax=axes[2], shrink=0.6)
+
+    # Plot 4: Grounding Attention Q2I
+    im3 = axes[3].imshow(grounding_attn_2d, cmap='inferno', aspect='auto')
+    axes[3].set_title('Grounding Attn')
+    axes[3].axis('off')
+    plt.colorbar(im3, ax=axes[3], shrink=0.6)
+
+    # Highlight sink tokens on the attention maps
+    if is_sink_token is not None:
+        # Find the coordinates of all sink tokens
+        sink_coords = np.argwhere(is_sink_token_2d)
+        for row, col in sink_coords:
+            # Define the properties of the highlight box
+            rect_args = {
+                'xy': (col - 0.5, row - 0.5), # Bottom-left corner of the cell
+                'width': 1,
+                'height': 1,
+                'linewidth': 2.5,
+                'edgecolor': '#7CFC00',  # A bright, "LawnGreen" color
+                'facecolor': 'none'
+            }
+            # Add a new rectangle patch to each of the three attention plots
+            #axes[1].add_patch(patches.Rectangle(**rect_args))
+            axes[2].add_patch(patches.Rectangle(**rect_args))
+            #axes[3].add_patch(patches.Rectangle(**rect_args))
+
+    if fg_bbox is not None:
+        fg_rect_arg = {
+            'xy': (fg_bbox[1] - 0.5, fg_bbox[0] - 0.5),  # Bottom-left corner of the cell
+            'width': fg_bbox[3] - fg_bbox[1] + 1,
+            'height': fg_bbox[2] - fg_bbox[0] + 1,
+            'linewidth': 2.5,
+            'edgecolor': '#FF4500',  # A bright, "OrangeRed" color
+            'facecolor': 'none'
+        }
+        axes[1].add_patch(patches.Rectangle(**fg_rect_arg))
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+
+    plt.show()
+    pass
+
+
+def save_attention_maps_separately(
+    image_path,
+    sink_attn,
+    grounding_attn_o2i,
+    grounding_attn,
+    save_dir,
+    is_sink_token=None,
+    fg_bbox=None
+):
+    """
+    Saves the original image and three attention maps as separate, clean image files
+    with no axes, titles, colorbars, or whitespace.
+
+    Args:
+        image_path (str): Path to the image file.
+        sink_attn (np.ndarray): Array (576,) for the first map (can be discrete or continuous).
+        grounding_attn_o2i (np.ndarray): Array (576,) for the second attention map.
+        grounding_attn (np.ndarray): Array (576,) for the third attention map.
+        save_dir (str): Directory where the output images will be saved.
+        is_sink_token (np.ndarray, optional): Boolean array (576,) to highlight sink tokens.
+        fg_bbox (list or tuple, optional): Bounding box [y_min, x_min, y_max, x_max] for the foreground.
+    """
+    # Ensure the save directory exists
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    # --- 1. Save Original Image ---
+    if os.path.exists(image_path):
+        img = Image.open(image_path)
+        img.save(os.path.join(save_dir, "original_image.png"))
+
+    # Reshape attention arrays and masks to 24x24
+    map1_data = sink_attn.reshape(h, w)
+    map2_data = grounding_attn_o2i.reshape(h, w)
+    map3_data = grounding_attn.reshape(h, w)
+
+    is_sink_token_2d = is_sink_token.reshape(h, w) if is_sink_token is not None else None
+
+    # --- 2. Process and Save the Three Attention Maps ---
+    map_data_list = [map1_data, map2_data, map3_data]
+    map_names = ["pseudo_label_or_attn_map", "grounding_attn_o2i", "grounding_attn"]
+    cmaps = ['viridis', 'inferno', 'inferno']  # Default cmaps
+
+    for i, (data, name) in enumerate(zip(map_data_list, map_names)):
+        fig, ax = plt.subplots(1, 1, figsize=(5, 5))
+        ax.axis('off')
+
+        # --- Plotting Logic ---
+        # Special handling for the first map if it's a discrete pseudo-label
+        if i == 0 and -1 in sink_attn:
+            colors = ['gray', 'black', '#FFE699']
+            cmap_discrete = ListedColormap(colors)
+            boundaries = [-1.5, -0.5, 0.5, 1.5]
+            norm = BoundaryNorm(boundaries, cmap_discrete.N)
+            ax.imshow(data, cmap=cmap_discrete, norm=norm, aspect='auto')
+        else:
+            ax.imshow(data, cmap=cmaps[i], aspect='auto')
+
+        # --- Overlay Logic ---
+        # Add sink token highlights (on all three maps)
+        if is_sink_token_2d is not None and i==1:
+            sink_coords = np.argwhere(is_sink_token_2d)
+            for row, col in sink_coords:
+                ax.add_patch(patches.Rectangle(
+                    (col - 0.5, row - 0.5), 1, 1,
+                    linewidth=2.5, edgecolor='#7CFC00', facecolor='none'
+                ))
+
+        # Add FG bounding box (only on the first map)
+        if i == 0 and fg_bbox is not None:
+            ax.add_patch(patches.Rectangle(
+                (fg_bbox[1] - 0.5, fg_bbox[0] - 0.5),
+                fg_bbox[3] - fg_bbox[1] + 1,
+                fg_bbox[2] - fg_bbox[0] + 1,
+                linewidth=2.5, edgecolor='#FF4500', facecolor='none'
+            ))
+
+        # --- Save the clean plot ---
+        save_path = os.path.join(save_dir, f"{name}.png")
+        fig.savefig(save_path, dpi=300, bbox_inches='tight', pad_inches=0)
+        plt.close(fig)  # Close the figure to free memory
+
+    print(f"Successfully saved 4 images to '{save_dir}'")
+
+
+
+def get_foreground_bbox(array_2d):
+    """
+    Get bounding box of foreground (non-zero) area in a 2D array.
+
+    Args:
+        array_2d: numpy array of shape (24, 24) or any 2D array
+
+    Returns:
+        tuple: (min_row, min_col, max_row, max_col) or None if no foreground
+               - min_row, min_col: top-left corner of bounding box
+               - max_row, max_col: bottom-right corner of bounding box (inclusive)
+    """
+    # Find all non-zero positions
+    fg_positions = np.where(array_2d > 0)
+    full_bbox = (0, 0, array_2d.shape[0] - 1, array_2d.shape[1] - 1)
+    # Check if there are any foreground pixels
+    if len(fg_positions[0]) == 0:
+        return full_bbox  # return full image as bbox if no foreground
+
+    # Get bounding box coordinates
+    min_row = np.min(fg_positions[0])
+    max_row = np.max(fg_positions[0])
+    min_col = np.min(fg_positions[1])
+    max_col = np.max(fg_positions[1])
+
+    return (min_row, min_col, max_row, max_col)
+
+def make_jsonl_file(json_path, save_path, save_name = 'llava_v1_5_selected_qa.jsonl', data_path='/home/yuheng/datasets', qa_pair='first'):
+    """
+    Save the json content to a jsonl file.
+    """
+    json_content = json.load(open(json_path, "r"))
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    save_file = os.path.join(save_path, save_name)
+    all_results = []
+    for idx, ele in enumerate(json_content):
+        if 'image' in ele.keys():# and 'coco' not in ele['image']:
+            question_id = str(idx)
+            image = os.path.join(data_path, ele['image'])
+            conversations = ele['conversations']
+            if qa_pair == 'first':
+                text = conversations[0]['value'].replace('<image>\n', '').replace('\n<image>', '')
+            else:
+                raise NotImplementedError(f"qa_pair={qa_pair} not implemented")
+            if '<image>' in text:
+                pass
+            result = {
+                "question_id": question_id,
+                "image": image,
+                "text": text,
+                "category": "default",
+            }
+            all_results.append(result)
+        else: continue
+    with open(save_file, "w") as f:
+        for result in all_results:
+            f.write(json.dumps(result) + "\n")
+    print(f"Saved {len(all_results)} results to {save_file}")
+
+def replace_img_path(pkl_path='/home/yuheng/datasets/llava_v1_5_pseudo_roi.pkl', data_path='/home/yuheng/datasets/'):
+    """
+    Replace the image path in the json file with the data path.
+    """
+    import pickle
+    pkl_content = pickle.load(open(pkl_path, "rb"))
+    all_results = []
+    for sample in pkl_content:
+        sample['image'] = sample['image'].replace(data_path, '')
+        all_results.append(sample)
+    dst_pkl_path = pkl_path.replace('.pkl', '_replaced.pkl')
+    with open(dst_pkl_path, "wb") as f:
+        pickle.dump(all_results, f)
+
+def replace_prompt(pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_3b_576res_v5.pkl',
+                   data_path='/home/yuheng/datasets/',
+                   dest_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_3b_576res_release.pkl'):
+    """
+    Replace the image path in the json file with the data path.
+    """
+    import pickle
+    pkl_content = pickle.load(open(pkl_path, "rb"))
+    all_results = []
+    for sample in pkl_content:
+        sample['prompt'] = sample['image'].replace('Output grounding bounding box related to the question in JSON.',
+                                    'Answer the question using a single word or phrase.')
+        all_results.append(sample)
+    if dest_pkl_path is None:
+        dst_pkl_path = pkl_path.replace('.pkl', '_replaced.pkl')
+    else:
+        dst_pkl_path = dest_pkl_path
+    with open(dst_pkl_path, "wb") as f:
+        pickle.dump(all_results, f)
+
+
+def refine_pseudo_labels_with_debug(original_image, pseudo_labels, std_threshold=5.0, grid_size=(24, 24), debug=True):
+    # 1. Prepare Image
+    if isinstance(original_image, Image.Image):
+        img_arr = np.array(original_image)
+    else:
+        img_arr = original_image
+
+    # Create Intensity Map (Max of RGB)
+    if len(img_arr.shape) == 3:
+        intensity_map = np.max(img_arr, axis=2)
+    else:
+        intensity_map = img_arr
+
+    H, W = intensity_map.shape
+    h_grid, w_grid = grid_size
+
+    # Reshape labels if flat
+    if pseudo_labels.shape != grid_size:
+        pseudo_labels = pseudo_labels.reshape(h_grid, w_grid)
+
+    refined_labels = pseudo_labels.copy()
+
+    step_h = H // h_grid
+    step_w = W // w_grid
+
+    rows, cols = np.where(pseudo_labels == 1)
+
+    # --- DEBUG COUNTER ---
+    debug_count = 0
+    MAX_DEBUG_PLOTS = 5  # Stop after showing this many plots
+
+    print(f"Checking {len(rows)} foreground tokens...")
+
+    for r, c in zip(rows, cols):
+        y_start = r * step_h
+        y_end = min((r + 1) * step_h, H)
+        x_start = c * step_w
+        x_end = min((c + 1) * step_w, W)
+
+        # Extract patches
+        intensity_patch = intensity_map[y_start:y_end, x_start:x_end]
+
+        # We also extract the RGB patch just for visualization
+        rgb_patch = img_arr[y_start:y_end, x_start:x_end]
+
+        if intensity_patch.size > 0:
+            patch_std = np.std(intensity_patch)
+
+            # --- VISUALIZATION BLOCK ---
+            if debug and debug_count < MAX_DEBUG_PLOTS:
+                fig, axes = plt.subplots(1, 2, figsize=(8, 4))
+
+                # Plot 1: Original RGB Patch (What the human sees)
+                axes[0].imshow(rgb_patch)
+                axes[0].set_title(f"RGB Patch\nGrid: ({r},{c})")
+                axes[0].axis('off')
+
+                # Plot 2: Intensity Patch (What the algo sees)
+                # Using 'jet' or 'magma' to visualize intensity better than gray
+                im = axes[1].imshow(intensity_patch, cmap='magma', vmin=0, vmax=255)
+                axes[1].set_title(f"Max-RGB Intensity\nStd Dev: {patch_std:.2f}")
+                axes[1].axis('off')
+
+                # Add status text
+                status = "KEEP" if patch_std >= std_threshold else "DROP (Too Flat)"
+                plt.suptitle(f"Token Analysis: {status} (Thresh: {std_threshold})",
+                             color='green' if patch_std >= std_threshold else 'red', weight='bold')
+
+                plt.tight_layout()
+                plt.show()
+
+                debug_count += 1
+            # ---------------------------
+
+            # 3. Apply Threshold
+            if patch_std < std_threshold:
+                refined_labels[r, c] = 0  # Revert to Background
+
+    return refined_labels
+def check_pseudo_label(dst_pkl_path, dataset_path = None):
+    """
+    Check the pseudo label for different datasets.
+    """
+    import pickle #10, 26, 62, 71, 102
+    from doctr.models import ocr_predictor
+    predictor = ocr_predictor(det_arch='fast_base', reco_arch='crnn_vgg16_bn', pretrained=True).cuda()
+    tmp_pkl_content = pickle.load(open(dst_pkl_path, "rb"))
+    for i in range(0, len(tmp_pkl_content)):
+        img_path = tmp_pkl_content[i]['image']
+        if dataset_path is not None: img_path = os.path.join(dataset_path, img_path)
+        #if 'ocr' not in img_path: continue
+        prompt = tmp_pkl_content[i]['prompt']
+        # if not ('left' in prompt or 'right' in prompt or 'top' in prompt or 'bottom' in prompt):
+        #     continue
+        response = tmp_pkl_content[i]['text']
+        print(f"Prompt: {prompt}\nResponse: {response}")
+        # if 'Is the bear sleeping?' not in prompt:
+        #     continue  # only check the specific prompt
+        grounding_attn_o2i = tmp_pkl_content[i]['grounding_attn_o2i']
+        sink_attn = tmp_pkl_content[i]['sink_attn']
+        # if 'Who is the author of this book?' not in prompt: continue
+        # if 'Suzanne Van Atten' not in tmp_pkl_content[i]['text']: continue
+        image = Image.open(img_path).convert('RGB')
+        image_size = image.size
+        known_fg_mask = get_foreground_mask(image_size, grounding_attn_o2i.shape[0], min_dimension_size=1)
+        tmp_attn = grounding_attn_o2i.copy()
+        tmp_attn[~known_fg_mask] = 0.
+        tmp_attn_binary = (tmp_attn > tmp_attn.max()*0.05).astype(np.float32)
+        # refined_labels = refine_pseudo_labels_with_debug(
+        #     original_image=expand2square(image, background_color=(127,127,127)),
+        #     pseudo_labels=tmp_attn_binary,
+        #     std_threshold=4.0,
+        #     grid_size=(24, 24),
+        #     debug=False
+        # )
+        refined_labels = refine_pseudo_labels_with_doctr(
+            original_image=expand2square(image, background_color=(127,127,127)),
+            pseudo_labels=tmp_attn_binary,
+            grid_size=(24, 24),
+            predictor=predictor,  # Pass the loaded model
+            debug=False
+        )
+
+        grounding_attn_o2i_refined = grounding_attn_o2i * refined_labels
+
+
+        # pseudo_set = create_pseudo_labels(sink_attn, grounding_attn_o2i_refined, sink_thresh=0.02, bg_coff=0.1, binary_coff=0.2,
+        #                                   original_image_size=image_size, for_vis=True, max_ratio_limit=0.5)
+        #
+        # plot_attention_analysis(
+        #     image_path=img_path,
+        #     sink_attn=pseudo_set['labels'],
+        #     grounding_attn_o2i=grounding_attn_o2i,
+        #     grounding_attn=pseudo_set['grounding_attn'],
+        #     is_sink_token=pseudo_set['sink_token_mask'],
+        #     fg_bbox=pseudo_set['fg_bbox'],
+        # )
+        #
+        pseudo_set = create_pseudo_labels(sink_attn, grounding_attn_o2i, sink_thresh=0.02, bg_coff=0.05, binary_coff=0.2,
+                                          original_image_size=image_size, for_vis=True, max_ratio_limit=0.5)
+
+
+        plot_attention_analysis(
+            image_path=img_path,
+            sink_attn=pseudo_set['labels'],
+            grounding_attn_o2i=grounding_attn_o2i,
+            grounding_attn=pseudo_set['grounding_attn'],
+            is_sink_token=pseudo_set['sink_token_mask'],
+            fg_bbox=pseudo_set['fg_bbox'],
+        )
+
+        # save_attention_maps_separately(
+        #     image_path=img_path,
+        #     sink_attn=pseudo_labels,
+        #     grounding_attn_o2i=grounding_attn_o2i,
+        #     grounding_attn=grounding_attn,
+        #     save_dir='/home/yuheng/code/llava/exclude/paper_plot',
+        #     is_sink_token=sink_tokens,
+        #     fg_bbox=pseudo_set['fg_bbox']
+        # )
+        pass
+
+    pass
+
+def refine_base_attn_with_grounding(grounding_pkl_path, base_pkl_path, dst_pkl_path, dataset_path):
+    """
+    Refine the base attention with grounding attention.
+    """
+    import pickle
+    grounding_pkl_content = pickle.load(open(grounding_pkl_path, "rb"))
+    base_pkl_content = pickle.load(open(base_pkl_path, "rb"))
+    assert len(grounding_pkl_content) == len(base_pkl_content), "The length of grounding and base pkl files should be the same."
+    all_results = []
+    sink_thresh, bg_coff, binary_coff = 0.02, 0.1, 0.2
+    for i in tqdm(range(0, len(grounding_pkl_content))):
+        assert grounding_pkl_content[i]['question_id'] == base_pkl_content[i]['question_id'], "The question_id should be the same."
+        assert grounding_pkl_content[i]['image'] == base_pkl_content[i]['image'], "The image path should be the same."
+        img_path = base_pkl_content[i]['image']
+        if dataset_path is not None: img_path = os.path.join(dataset_path, img_path)
+        image = Image.open(img_path).convert('RGB')
+        image_size = image.size
+        grounding_sample = grounding_pkl_content[i]
+        grounding_pseudo_set = create_pseudo_labels(grounding_sample['sink_attn'],
+                                          grounding_sample['grounding_attn_o2i'],
+                                          sink_thresh=sink_thresh,
+                                          bg_coff=bg_coff,
+                                          binary_coff=binary_coff,
+                                          original_image_size=image_size,
+                                          for_vis=True,
+                                          max_ratio_limit=0.5,
+                                          use_smoothing=True
+                                        )
+        grounding_bg_mask = grounding_pseudo_set['labels'] == 0  # background mask
+        # plot_attention_analysis(
+        #     image_path=img_path,
+        #     sink_attn=grounding_pseudo_set['labels'],
+        #     grounding_attn_o2i=grounding_sample['grounding_attn_o2i'],
+        #     grounding_attn=grounding_pseudo_set['grounding_attn'],
+        #     is_sink_token=grounding_pseudo_set['sink_token_mask'],
+        #     fg_bbox=grounding_pseudo_set['fg_bbox'],
+        # )
+
+        # Refine the base attention with grounding background mask
+        base_sample = base_pkl_content[i]
+        base_grounding_attn_o2i = base_sample['grounding_attn_o2i']
+        feat_shape = base_grounding_attn_o2i.shape
+        base_grounding_attn_o2i = base_grounding_attn_o2i.flatten()
+        base_grounding_attn_o2i[grounding_bg_mask.flatten()] = 0.
+        base_grounding_attn_o2i = base_grounding_attn_o2i.reshape(feat_shape)
+        # base_pseudo_set = create_pseudo_labels(base_sample['sink_attn'],
+        #                                     base_grounding_attn_o2i,
+        #                                     sink_thresh=sink_thresh,
+        #                                     bg_coff=bg_coff,
+        #                                     binary_coff=binary_coff,
+        #                                     original_image_size=image_size,
+        #                                     for_vis=True,
+        #                                     max_ratio_limit=0.5)
+        # plot_attention_analysis(
+        #     image_path=img_path,
+        #     sink_attn=base_pseudo_set['labels'],
+        #     grounding_attn_o2i=base_grounding_attn_o2i,
+        #     grounding_attn=base_pseudo_set['grounding_attn'],
+        #     is_sink_token=base_pseudo_set['sink_token_mask'],
+        #     fg_bbox=base_pseudo_set['fg_bbox'],
+        # )
+
+        refined_sample = {
+            "question_id": base_sample['question_id'],
+            "prompt": base_sample['prompt'],
+            "text": base_sample['text'],
+            "image": base_sample['image'],
+            "sink_attn": base_sample['sink_attn'],
+            "grounding_attn_o2i": base_grounding_attn_o2i
+        }
+
+        all_results.append(refined_sample)
+        pass
+
+    with open(dst_pkl_path, "wb") as f:
+        pickle.dump(all_results, f)
+    print(f"Saved {len(all_results)} results to {dst_pkl_path}")
+    return dst_pkl_path
+
+def merge_pkls(pkl_list, dst_path, dst_pkl_name='llava_v1_5_pseudo_roi.pkl'):
+    """
+    Get the pkl files from the list and save them to a single pkl file.
+    """
+    import pickle
+    all_results = []
+    for pkl_path in pkl_list:
+        tmp_pkl_content = pickle.load(open(pkl_path, "rb"))
+        all_results.extend(tmp_pkl_content)
+    os.makedirs(dst_path, exist_ok=True)
+    dst_pkl_path = os.path.join(dst_path, dst_pkl_name)
+    with open(dst_pkl_path, "wb") as f:
+        pickle.dump(all_results, f)
+    print(f"Saved {len(all_results)} results to {dst_pkl_path}")
+    return dst_pkl_path
+
+
+# dataset_dir = '/home/yuheng/datasets/'
+# dst_json_path = os.path.join(dataset_dir,'llava_v1_5_mix665k.json') #"/home/yuheng/datasets/llava_v1_5_mix665k.json"
+# tmp_json_content = json.load(open(dst_json_path, "r"))
+# json_content_nococo =  [ele for ele in tmp_json_content if 'image' in ele.keys() and 'coco' not in ele['image']]
+# coco = [ele for ele in tmp_json_content if 'image' in ele.keys() and 'coco' in ele['image']]
+# vg_100k = [ele for ele in json_content_nococo if 'image' in ele.keys() and 'VG_100K' in ele['image']]
+# gqa = [ele for ele in json_content_nococo if 'image' in ele.keys() and 'gqa' in ele['image']]
+# textvqa = [ele for ele in json_content_nococo if 'image' in ele.keys() and 'textvqa' in ele['image']]
+# ocr_vqa = [ele for ele in json_content_nococo if 'image' in ele.keys() and 'ocr_vqa' in ele['image']]
+# pass
+#
+#
+# test_json_152k = os.path.join(dataset_dir,'llava_v1_5_mix152k_test.json') #'/home/yuheng/datasets/llava_v1_5_mix152k_test.json'
+# test_json_152k_content = gqa + ocr_vqa
+# with open(test_json_152k, "w") as f:
+#     json.dump(test_json_152k_content, f, indent=2)
+
+# test_json_238k = os.path.join(dataset_dir,'llava_v1_5_mix238k_test.json') #'/home/yuheng/datasets/llava_v1_5_mix152k_test.json'
+# test_json_238k_content = gqa + ocr_vqa + vg_100k
+# with open(test_json_238k, "w") as f:
+#     json.dump(test_json_238k_content, f, indent=2)
+
+
+
+# make_jsonl_file(json_path='/home/yuheng/datasets/llava_v1_5_mix665k.json',
+#                  save_path='/home/yuheng/datasets',
+#                  save_name='llava_v1_5_mix665k_selected_qa.jsonl',
+#                  data_path='/home/yuheng/datasets',
+#                  qa_pair='first')
+
+if __name__ == "__main__":
+    # make_jsonl_file(json_path='/mnt/sda/yuheng/datasets/llava_v1_5_mix665k.json',
+    #                  save_path='/mnt/sda/yuheng/datasets',
+    #                  save_name='llava_v1_5_mix665k_selected_qa.jsonl',
+    #                  data_path='/mnt/sda/yuheng/datasets',
+    #                  qa_pair='first')
+    #replace_img_path('/home/yuheng/datasets/llava_v1_5_pseudo_roi_replaced.pkl')
+    #check_pseudo_label(dst_pkl_path='/mnt/sda/yuheng/datasets/llava_v1_5_pseudo_roi_replaced.pkl',  dataset_path='/mnt/sda/yuheng/datasets/')
+    #check_pseudo_label(dst_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_576res_229k_denoised.pkl', dataset_path='/home/yuheng/datasets/')
+    # check_pseudo_label(dst_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_docvqa33k_fullPromptFix_denoised.pkl', dataset_path='/home/yuheng/datasets/')
+
+    # refine_base_attn_with_grounding(
+    #     grounding_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_3b_gqa_576res_25k.pkl',
+    #     base_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_3b_gqa_576res_v3_25k.pkl',
+    #     dst_pkl_path='/home/yuheng/datasets/qwen_v2_5_pseudo_3b_gqa_refined.pkl',
+    #     dataset_path='/home/yuheng/datasets/'
+    # )
+    # replace_prompt()
+    # pass
+
+    merge_pkls(    pkl_list=[
+            '/root/epfs/sproject/datasets/qwen2_5vl_pseudo_7b_576res_release_denoisedV2.pkl',
+            '/root/epfs/sproject/datasets/qwen2_5vl_pseudo_7b_docvqa33k_fullPromptFix_denoised.pkl'
+        ],
+        dst_path='/root/epfs/sproject/datasets',
+        dst_pkl_name='qwen2_5vl_pseudo_7b_576res_185k_denoisedV2.pkl')
+    # # replace_img_path('/home/yuheng/datasets/llava_v1_5_13b_pseudo_roi.pkl')
+
+
+# zero_array = np.zeros(576)
+# plot_attention_analysis(os.path.join(dataset_dir,'gqa/images/2348106.jpg'),zero_array,zero_array,zero_array)
+# pass
+#{'from': 'human', 'value': 'Please provide the bounding box coordinate of the region this sentence describes: A black corner lamp.'}
+

QWENVL-PRIVATE/analysis/check_textvqa_json.py

@@ -0,0 +1,56 @@
+import json
+
+# --- 1. Load Source Data ---
+# (Using your provided paths)
+# textvqa_train_json = '/home/yuheng/datasets/textvqa/TextVQA_0.5.1_train.json'
+# output_jsonl_file = '/home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl'
+
+textvqa_train_json = '/home/yuheng/code/llava/playground/data/eval/textvqa/TextVQA_0.5.1_val.json'
+output_jsonl_file = '/home/yuheng/datasets/textvqa/converted_llava_style_textvqa_val.jsonl'
+
+try:
+    # Load the entire JSON file
+    textvqa_train_data = json.load(open(textvqa_train_json))['data']
+
+    print(f"Loaded {len(textvqa_train_data)} items from {textvqa_train_json}")
+
+    converted_data = []
+
+    # --- 2. Iterate and Convert Each Item ---
+    for item in textvqa_train_data:
+        # Get the image_id (string) from the source
+        # This is used for both 'question_id' and 'image' in the target format
+        image_id = item.get('image_id')
+
+        # Get the raw question text
+        question_text = item.get('question')
+
+        # Create the new dictionary in the target format
+        new_item = {
+            "question_id": image_id,
+            "image": f"{image_id}.jpg",
+            "text": question_text,  # See critical note below
+            "category": "default",  # Hardcoded, as in your example
+            "answers": item.get('answers', [])  # List of answers
+        }
+
+        converted_data.append(new_item)
+
+    # --- 3. Write to JSONL (JSON Lines) File ---
+    # The target format is a .jsonl file, where each line is a separate JSON object.
+    with open(output_jsonl_file, 'w') as f:
+        for entry in converted_data:
+            # Dump each dictionary as a string and add a newline
+            f.write(json.dumps(entry) + '\n')
+
+    print(f"Successfully converted and saved to {output_jsonl_file}")
+
+    # Print a sample of the first converted item
+    if converted_data:
+        print("\n--- Sample of first converted item ---")
+        print(json.dumps(converted_data[0], indent=2))
+
+except FileNotFoundError:
+    print(f"Error: Source file not found at {textvqa_train_json}")
+except Exception as e:
+    print(f"An error occurred: {e}")

QWENVL-PRIVATE/analysis/convert_chartqa.py

@@ -0,0 +1,84 @@
+import os
+import json
+from datasets import load_dataset
+from tqdm import tqdm
+
+
+def process_chartqa_dataset(base_dataset_dir, output_image_dir, output_jsonl_path):
+    """
+    Downloads ChartQA, saves images, and writes JSONL with paths relative
+    to base_dataset_dir.
+    """
+
+    # 1. Create the output directory if it doesn't exist
+    if not os.path.exists(output_image_dir):
+        os.makedirs(output_image_dir)
+        print(f"Created directory: {output_image_dir}")
+
+    # 2. Load the dataset
+    print("Loading HuggingFaceM4/ChartQA dataset...")
+    try:
+        dataset = load_dataset("HuggingFaceM4/ChartQA", split="train")
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return
+
+    print(f"Processing {len(dataset)} items...")
+
+    # 3. Iterate and Process
+    with open(output_jsonl_path, 'w', encoding='utf-8') as f:
+        for idx, item in enumerate(tqdm(dataset)):
+            try:
+                # --- A. Save Image ---
+                image = item['image']
+
+                # Create filename
+                image_filename = f"chartqa_train_{idx}.png"
+                full_image_path = os.path.join(output_image_dir, image_filename)
+
+                # Save the image physically to the full path
+                image.save(full_image_path)
+
+                # --- B. Calculate Relative Path ---
+                # This converts /home/yuheng/datasets/ChartQA/images/img.png
+                # to ChartQA/images/img.png
+                relative_image_path = os.path.relpath(full_image_path, base_dataset_dir)
+
+                # --- C. Process Data Fields ---
+                question_text = item.get('query', item.get('question', ''))
+                raw_label = item.get('label', [])
+                answers_list = raw_label if isinstance(raw_label, list) else [str(raw_label)]
+                primary_answer = answers_list[0] if answers_list else ""
+
+                # --- D. Construct JSON Object ---
+                entry = {
+                    "question_id": str(idx),
+                    "image": relative_image_path,  # <--- UPDATED: Uses relative path now
+                    "prompt": question_text,
+                    "category": "default",
+                    "answers": answers_list,
+                    "answer": primary_answer,
+                    "dataset": "ChartQA"
+                }
+
+                f.write(json.dumps(entry) + "\n")
+
+            except Exception as e:
+                print(f"Skipping index {idx} due to error: {e}")
+
+    print(f"\nProcessing complete.")
+    print(f"Images saved to: {output_image_dir}")
+    print(f"Annotations saved to: {output_jsonl_path}")
+
+
+# --- Configuration ---
+if __name__ == "__main__":
+    # The root directory you want the paths to be relative to
+    BASE_ROOT = "/home/yuheng/datasets"
+
+    # Where the images and jsonl will actually be stored
+    # Ensure this path is inside BASE_ROOT if you want a clean relative path
+    IMAGE_SAVE_PATH = os.path.join(BASE_ROOT, "ChartQA/images")
+    JSONL_SAVE_PATH = os.path.join(BASE_ROOT, "chartqa_28k_qa.jsonl")
+
+    process_chartqa_dataset(BASE_ROOT, IMAGE_SAVE_PATH, JSONL_SAVE_PATH)

QWENVL-PRIVATE/analysis/convert_json2qwen_training.py

@@ -0,0 +1,220 @@
+import json
+import argparse
+import os
+import random
+import glob
+
+
+def fix_image_path(raw_image_id, data_root):
+    """
+    Tries to find the image in data_root.
+    If not found, applies dataset-specific folder corrections (like textvqa).
+    """
+    full_path = os.path.join(data_root, raw_image_id)
+    if os.path.exists(full_path):
+        return raw_image_id
+
+    # Heuristic fixes based on your previous logs
+    # Check if it needs 'textvqa/train_images' prefix
+    if "textvqa" in raw_image_id.lower() or len(raw_image_id) == 16:  # e.g. 20de8a999becc896.jpg
+        # Try appending textvqa prefix if not present
+        potential_fix = os.path.join('textvqa/train_images', os.path.basename(raw_image_id))
+        if os.path.exists(os.path.join(data_root, potential_fix)):
+            print(f"  Fixed path: {raw_image_id} -> {potential_fix}")
+            return potential_fix
+
+    # Add other dataset fixes here (DocVQA, ChartQA, etc.)
+    # Example for DocVQA:
+    # if "docvqa" in raw_image_id.lower(): ...
+
+    return raw_image_id
+
+
+def convert_to_qwen_mixed(
+        reg_file, imp_file, output_file, data_root,
+        mixing_ratio=1.0, use_gt=False, seed=42,
+        use_short_prompt=False
+):
+    random.seed(seed)
+
+    # 1. Load Data
+    print(f"Loading Regressions: {reg_file}")
+    with open(reg_file, 'r', encoding='utf-8') as f:
+        reg_data = json.load(f)
+
+    imp_data = []
+    if imp_file and os.path.exists(imp_file):
+        print(f"Loading Improvements: {imp_file}")
+        with open(imp_file, 'r', encoding='utf-8') as f:
+            imp_data = json.load(f)
+    else:
+        print("No improvements file provided or found. Skipping mixing.")
+
+    # 2. Determine Counts based on Ratio
+    num_regs = len(reg_data)
+    target_imp_count = int(num_regs * mixing_ratio)
+
+    print(f"stat: {num_regs} Regressions found.")
+
+    selected_improvements = []
+    if imp_data:
+        if len(imp_data) > target_imp_count:
+            print(f"Sampling {target_imp_count} Improvements from {len(imp_data)} available (Ratio {mixing_ratio})")
+            selected_improvements = random.sample(imp_data, target_imp_count)
+        else:
+            print(f"Taking all {len(imp_data)} available Improvements (Target was {target_imp_count})")
+            selected_improvements = imp_data
+
+    # 3. Tag Sources and Assign Targets
+    # Regressions -> Target is Baseline (or GT)
+    # Improvements -> Target is ROI (or GT)
+
+    final_pool = []
+
+    # Process Regressions
+    for item in reg_data:
+        item['source_type'] = 'regression'
+        # For regression, Baseline is correct
+        item['target_response_key'] = 'baseline_response'
+        final_pool.append(item)
+
+    # Process Improvements
+    for item in selected_improvements:
+        item['source_type'] = 'improvement'
+        # For improvement, ROI is correct (Baseline is wrong!)
+        item['target_response_key'] = 'roi_response'
+        final_pool.append(item)
+
+    # Shuffle the pool
+    random.shuffle(final_pool)
+
+    print(f"Generating {len(final_pool)} training samples...")
+
+    # 4. Conversion Loop
+    with open(output_file, 'w', encoding='utf-8') as f_out:
+        for item in final_pool:
+            # A. Select Response
+            if use_gt:
+                # If forcing GT, ignore model responses
+                gt = item.get('ground_truth', '')
+                if isinstance(gt, list):
+                    target_response = str(gt[0])
+                else:
+                    target_response = str(gt)
+            else:
+                # Use the key determined above (Base for Reg, ROI for Imp)
+                key = item['target_response_key']
+                target_response = item.get(key, '')
+
+            # B. Format Question
+            question_text = item.get('question', '')
+            if "<image>" not in question_text:
+                question_text = "<image>\n" + question_text
+
+            if use_short_prompt and item['source_type']== 'improvement':
+                question_text = question_text + "\nAnswer the question using a single word or phrase."
+
+            # C. Fix Image Path
+            raw_image_id = item.get('image_id', '')
+            if not os.path.exists(os.path.join(data_root, raw_image_id)):
+                raw_image_id = os.path.join('textvqa/train_images', raw_image_id)
+                print("Adjusted image path to:", raw_image_id)
+            fixed_image_id = raw_image_id
+
+            # Optional: Warning if still missing
+            if not os.path.exists(os.path.join(data_root, fixed_image_id)):
+                # Just a warning, don't crash, maybe the path is absolute or handled elsewhere
+                print(f"  Warning: Image not found at {os.path.join(data_root, fixed_image_id)}")
+                pass
+
+                # D. Construct Entry
+            entry = {
+                "id": str(item.get('question_id', 'unknown')),
+                "image": fixed_image_id,
+                "data_path": data_root,
+                "conversations": [
+                    {
+                        "from": "human",
+                        "value": question_text
+                    },
+                    {
+                        "from": "gpt",
+                        "value": target_response
+                    }
+                ]
+            }
+
+            f_out.write(json.dumps(entry) + '\n')
+
+    print(f"✅ Conversion complete.")
+    print(f"Saved to '{output_file}'")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert Mixed Reg/Imp to Qwen Training Format")
+    parser.add_argument('--reg-file', type=str, required=True, help="Path to regressions.json")
+    parser.add_argument('--imp-file', type=str, default=None, help="Path to improvements.json")
+    parser.add_argument('--output', type=str, default="finetune_mixed.jsonl", help="Output JSONL path")
+    parser.add_argument('--data-root', type=str, default="./data", help="Root path for images")
+    parser.add_argument('--mixing-ratio', type=float, default=1.0,
+                        help="Ratio of Improvements to keep relative to Regressions (e.g. 0.5 = 1 Imp for 2 Regs)")
+    parser.add_argument('--use-gt', action='store_true', help="Force using Ground Truth for all answers")
+    parser.add_argument('--use-short-prompt', action='store_true', help="Whether to use short prompt format")
+    args = parser.parse_args()
+
+    if os.path.exists(args.reg_file):
+        convert_to_qwen_mixed(
+            args.reg_file, args.imp_file, args.output,
+            args.data_root, args.mixing_ratio, args.use_gt,
+            use_short_prompt=args.use_short_prompt
+        )
+    else:
+        print(f"Error: Regression file '{args.reg_file}' not found.")
+    '''
+    python analysis/convert_json2qwen_training.py --input /home/yuheng/datasets/regressions-v2.json --output /home/yuheng/datasets/roi-post-sft7k.jsonl --data-root /home/yuheng/datasets/
+    python analysis/convert_json2qwen_training.py --input /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2.json --output /root/epfs/sproject/datasets/roi-post-sft7k.jsonl --data-root /root/epfs/sproject/datasets
+    
+    
+    python analysis/convert_json2qwen_training.py \
+    --reg-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2.json \
+    --imp-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2.json \
+    --output /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/mixed_finetune_v2.jsonl \
+    --data-root /root/epfs/sproject/datasets \
+    --mixing-ratio 1.0
+    
+    
+    
+    python analysis/convert_json2qwen_training.py \
+    --reg-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2-short.json \
+    --imp-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-short.json \
+    --output /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/mixed_finetune_v2-short.jsonl \
+    --data-root /root/epfs/sproject/datasets \
+    --use-short-prompt \
+    --mixing-ratio 1.0
+    
+    python analysis/convert_json2qwen_training.py \
+    --reg-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2-short.json \
+    --imp-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-short.json \
+    --output /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/roi-post-short-sft4k.jsonl \
+    --data-root /root/epfs/sproject/datasets \
+    --use-short-prompt \
+    --mixing-ratio 0.0
+    
+    
+    python analysis/convert_json2qwen_training.py \
+    --reg-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2.json \
+    --imp-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-short.json \
+    --output /root/epfs/sproject/datasets/roi-post-mixed-sft14k.jsonl \
+    --data-root /root/epfs/sproject/datasets \
+    --use-short-prompt \
+    --mixing-ratio 1.0
+    
+    
+    python analysis/convert_json2qwen_training.py \
+    --reg-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-clean-32b-judge-wo-spatial.json  \
+    --imp-file /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2-32b-judge.json \
+    --output /root/epfs/sproject/datasets/roi-post-mixed-sft12k.jsonl \
+    --data-root /root/epfs/sproject/datasets \
+    --mixing-ratio 1.0
+    
+    '''

QWENVL-PRIVATE/analysis/convert_seal_json.py

@@ -0,0 +1,81 @@
+import json
+import re
+import os
+
+
+def clean_prompt(text):
+    # Matches: <image> followed by newlines, then "Additional visual info..." ending with a dot and newline
+    pattern = r"<image>\s*\n\s*Additional visual information to focus on:.*?\.\n"
+    cleaned_text = re.sub(pattern, "", text, flags=re.DOTALL)
+    return cleaned_text.strip()
+
+
+def convert_json_to_jsonl(input_file, output_file):
+    with open(input_file, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+
+    with open(output_file, 'w', encoding='utf-8') as f_out:
+        total_extracted = 0
+        skipped_count = 0
+
+        for item_idx, item in enumerate(data):
+            original_path = item.get('image', '')
+
+            # 1. FILTER: Only process if it starts with 'coco2017'
+            if not original_path.startswith('coco2017'):
+                skipped_count += 1
+                continue
+
+            # 2. RENAME: Replace only the first occurrence of 'coco2017' with 'coco'
+            # e.g., "coco2017/train2017/..." -> "coco/train2017/..."
+            new_image_path = original_path.replace('coco2017', 'coco', 1)
+
+            # Generate ID
+            base_id = os.path.splitext(os.path.basename(original_path))[0]
+            if not base_id:
+                base_id = f"item_{item_idx}"
+
+            conversations = item.get('conversations', [])
+
+            # Process turns
+            for i in range(0, len(conversations), 2):
+                if i + 1 >= len(conversations):
+                    break
+
+                human_turn = conversations[i]
+                gpt_turn = conversations[i + 1]
+
+                if human_turn['from'] != 'human' or gpt_turn['from'] != 'gpt':
+                    continue
+
+                question_text = clean_prompt(human_turn['value'])
+                answer_text = gpt_turn['value']
+                q_id = f"{base_id}_{i // 2}"
+
+                new_entry = {
+                    "question_id": q_id,
+                    "image": new_image_path,  # Uses the modified 'coco' path
+                    "text": question_text,
+                    "answers": [answer_text],
+                    "category": "default"
+                }
+
+                f_out.write(json.dumps(new_entry) + '\n')
+                total_extracted += 1
+
+    print(f"Conversion complete.")
+    print(f"Skipped {skipped_count} items (not from coco2017).")
+    print(f"Generated {total_extracted} JSONL lines in '{output_file}'.")
+
+
+# --- Usage ---
+if __name__ == "__main__":
+    # Example usage
+    # Ensure you change these filenames to your actual paths
+    input_json = "/home/yuheng/datasets/spatial_relation_data.json"
+    output_jsonl = "/home/yuheng/datasets/vstar_coco_spatial_relation_data.jsonl"
+
+    if os.path.exists(input_json):
+        convert_json_to_jsonl(input_json, output_jsonl)
+    else:
+        print(f"File {input_json} not found. Please set the correct path.")

QWENVL-PRIVATE/analysis/denoise_textual_file.py

@@ -0,0 +1,348 @@
+import sys
+import os
+
+# Add the parent directory to sys.path to import qwen_src
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import pickle
+import numpy as np
+from PIL import Image
+from multiprocessing import Pool, cpu_count
+from functools import partial
+from tqdm import tqdm
+import matplotlib.pyplot as plt  # Needed for the debug function if enabled
+import cv2
+
+# Assuming these are available in your environment
+from qwen_src.mm_utils import expand2square, create_pseudo_labels, get_foreground_mask
+
+
+# Import Doctr
+from doctr.models import ocr_predictor
+
+# --- GLOBAL LOADER TO PREVENT RE-LOADING MODEL EVERY CALL ---
+# In a production loop, initialize this ONCE outside your loop and pass it in.
+_GLOBAL_PREDICTOR = None
+
+def get_doctr_predictor():
+    global _GLOBAL_PREDICTOR
+    if _GLOBAL_PREDICTOR is None:
+        # We only need detection (det), recognition (reco) is not needed for masking
+        # causing a 2x speedup.
+        print("Loading docTR model (one-time setup)...")
+        _GLOBAL_PREDICTOR = ocr_predictor(det_arch='db_resnet50', reco_arch=None, pretrained=True)
+    return _GLOBAL_PREDICTOR
+
+
+def refine_pseudo_labels_with_doctr(original_image, pseudo_labels, grid_size=(24, 24), overlap_threshold=0.15,
+                                    debug=False, predictor=None):
+    """
+    Refines pseudo-labels by intersecting them with text regions detected by docTR.
+
+    Args:
+        original_image: PIL Image or Numpy Array.
+        pseudo_labels: (H_grid, W_grid) binary mask.
+        grid_size: tuple (h_grid, w_grid).
+        overlap_threshold: float (0.0 to 1.0). If a grid token contains this percentage
+                           of text pixels, it is kept.
+        predictor: Optional passed instance of ocr_predictor to save loading time.
+    """
+
+    # 1. Prepare Image
+    if isinstance(original_image, Image.Image):
+        img_arr = np.array(original_image)
+    else:
+        img_arr = original_image.copy()
+
+    H, W = img_arr.shape[:2]
+
+    # 2. Get Predictor
+    if predictor is None:
+        doc_model = get_doctr_predictor()
+    else:
+        doc_model = predictor
+
+    # 3. Run Text Detection (Inference)
+    # docTR expects list of numpy arrays
+    result = doc_model([img_arr])
+
+    # 4. Generate a "Text Mask" at full resolution
+    # We will draw white boxes on a black canvas where text is found.
+    text_mask = np.zeros((H, W), dtype=np.uint8)
+
+    # Extract boxes. Doctr returns relative coords (xmin, ymin, xmax, ymax) 0.0-1.0
+    # Structure: result.pages[0].blocks[i].lines[j].words[k]
+    # We use 'words' for tightest fit, or 'lines' for better connectivity.
+    # Let's use words to be precise.
+
+    boxes = []
+    for block in result.pages[0].blocks:
+        for line in block.lines:
+            for word in line.words:
+                # geometry is ((xmin, ymin), (xmax, ymax))
+                (min_x, min_y), (max_x, max_y) = word.geometry
+
+                # Convert to absolute pixels
+                abs_xmin = int(min_x * W)
+                abs_xmax = int(max_x * W)
+                abs_ymin = int(min_y * H)
+                abs_ymax = int(max_y * H)
+
+                boxes.append([abs_xmin, abs_ymin, abs_xmax, abs_ymax])
+
+                # Draw filled rectangle on the mask
+                cv2.rectangle(text_mask, (abs_xmin, abs_ymin), (abs_xmax, abs_ymax), 1, thickness=-1)
+
+    # 5. Intersect Grid with Text Mask
+    h_grid, w_grid = grid_size
+
+    # Ensure labels shape
+    if pseudo_labels.shape != grid_size:
+        pseudo_labels = pseudo_labels.reshape(h_grid, w_grid)
+
+    refined_labels = pseudo_labels.copy()
+    step_h = H // h_grid
+    step_w = W // w_grid
+
+    rows, cols = np.where(pseudo_labels == 1)
+
+    debug_count = 0
+    MAX_DEBUG_PLOTS = 3
+
+    for r, c in zip(rows, cols):
+        y_start = r * step_h
+        y_end = min((r + 1) * step_h, H)
+        x_start = c * step_w
+        x_end = min((c + 1) * step_w, W)
+
+        # Slice the text mask corresponding to this grid token
+        mask_patch = text_mask[y_start:y_end, x_start:x_end]
+
+        # Calculate Intersection: How much of this patch is actually text?
+        pixel_count = mask_patch.size
+        text_pixel_count = np.sum(mask_patch)
+
+        ratio = text_pixel_count / pixel_count if pixel_count > 0 else 0
+
+        # Logic: If the Grid Token (Attention) is ON, but there is NO text in that region, turn it OFF.
+        keep_token = ratio >= overlap_threshold
+
+        if not keep_token:
+            refined_labels[r, c] = 0
+
+        # --- DEBUG VISUALIZATION ---
+        if debug and debug_count < MAX_DEBUG_PLOTS:
+            # Only show borderline cases or drops to save time
+            # visualizing patches that were 1 but dropped to 0
+            if not keep_token:
+                fig, axes = plt.subplots(1, 3, figsize=(12, 4))
+
+                # RGB
+                axes[0].imshow(img_arr[y_start:y_end, x_start:x_end])
+                axes[0].set_title(f"RGB Patch ({r},{c})")
+
+                # Global Text Mask Context
+                # Show the whole mask with a red box around current patch
+                debug_mask = text_mask.copy() * 255
+                cv2.rectangle(debug_mask, (x_start, y_start), (x_end, y_end), 127, 2)
+                axes[1].imshow(debug_mask, cmap='gray')
+                axes[1].set_title("Global Text Detections")
+
+                # Local Intersection
+                axes[2].imshow(mask_patch, vmin=0, vmax=1)
+                axes[2].set_title(f"Overlap Ratio: {ratio:.2f}\nThresh: {overlap_threshold}")
+
+                plt.suptitle(f"Action: {'KEEP' if keep_token else 'DROP'}",
+                             color='green' if keep_token else 'red', weight='bold')
+                plt.show()
+                debug_count += 1
+        # ---------------------------
+
+    return refined_labels
+
+
+# --- Your existing refine function (kept as is) ---
+def refine_pseudo_labels_with_debug(original_image, pseudo_labels, std_threshold=5.0, grid_size=(24, 24), debug=True):
+    # 1. Prepare Image
+    if isinstance(original_image, Image.Image):
+        img_arr = np.array(original_image)
+    else:
+        img_arr = original_image
+
+    # Create Intensity Map (Max of RGB)
+    if len(img_arr.shape) == 3:
+        intensity_map = np.max(img_arr, axis=2)
+    else:
+        intensity_map = img_arr
+
+    H, W = intensity_map.shape
+    h_grid, w_grid = grid_size
+
+    # Reshape labels if flat
+    if pseudo_labels.shape != grid_size:
+        pseudo_labels = pseudo_labels.reshape(h_grid, w_grid)
+
+    refined_labels = pseudo_labels.copy()
+
+    step_h = H // h_grid
+    step_w = W // w_grid
+
+    rows, cols = np.where(pseudo_labels == 1)
+
+    # --- DEBUG COUNTER ---
+    debug_count = 0
+    MAX_DEBUG_PLOTS = 5
+
+    for r, c in zip(rows, cols):
+        y_start = r * step_h
+        y_end = min((r + 1) * step_h, H)
+        x_start = c * step_w
+        x_end = min((c + 1) * step_w, W)
+
+        # Extract patches
+        intensity_patch = intensity_map[y_start:y_end, x_start:x_end]
+        rgb_patch = img_arr[y_start:y_end, x_start:x_end]
+
+        if intensity_patch.size > 0:
+            patch_std = np.std(intensity_patch)
+
+            # --- VISUALIZATION BLOCK ---
+            # Note: Debug plotting inside multiprocessing usually fails or is hidden.
+            # Ensure debug=False when running in parallel.
+            if debug and debug_count < MAX_DEBUG_PLOTS:
+                fig, axes = plt.subplots(1, 2, figsize=(8, 4))
+                axes[0].imshow(rgb_patch)
+                axes[0].set_title(f"RGB Patch\nGrid: ({r},{c})")
+                axes[0].axis('off')
+                im = axes[1].imshow(intensity_patch, cmap='magma', vmin=0, vmax=255)
+                axes[1].set_title(f"Max-RGB Intensity\nStd Dev: {patch_std:.2f}")
+                axes[1].axis('off')
+                status = "KEEP" if patch_std >= std_threshold else "DROP (Too Flat)"
+                plt.suptitle(f"Token Analysis: {status} (Thresh: {std_threshold})",
+                             color='green' if patch_std >= std_threshold else 'red', weight='bold')
+                plt.tight_layout()
+                plt.show()
+                debug_count += 1
+            # ---------------------------
+
+            if patch_std < std_threshold:
+                refined_labels[r, c] = 0
+
+    return refined_labels
+
+
+# --- Worker Function for Multiprocessing ---
+def process_single_sample(item, dataset_path):
+    """
+    Worker function to process a single data item.
+    Returns the modified item or None if processing fails.
+    """
+    try:
+        img_path = item['image']
+        if 'gqa' in img_path.lower():
+            return item # Skip GQA images
+        if dataset_path is not None:
+            full_img_path = os.path.join(dataset_path, img_path)
+        else:
+            full_img_path = img_path
+
+        grounding_attn_o2i = item['grounding_attn_o2i']
+
+        # Load and convert image
+        try:
+            image = Image.open(full_img_path).convert('RGB')
+        except Exception as e:
+            # print(f"Error opening image {full_img_path}: {e}")
+            return None
+
+        image_size = image.size
+
+        # Calculate mask
+        known_fg_mask = get_foreground_mask(image_size, grounding_attn_o2i.shape[0], min_dimension_size=1)
+        tmp_attn = grounding_attn_o2i.copy()
+        tmp_attn[~known_fg_mask] = 0.
+
+        # Binary Threshold
+        tmp_attn_binary = (tmp_attn > tmp_attn.max() * 0.05).astype(np.float32)
+
+        # Denoise (Force debug=False for multiprocessing)
+        refined_labels = refine_pseudo_labels_with_debug(
+            original_image=expand2square(image, background_color=(127, 127, 127)),
+            pseudo_labels=tmp_attn_binary,
+            std_threshold=4.0,
+            grid_size=(grounding_attn_o2i.shape[0], grounding_attn_o2i.shape[1]),
+            debug=False
+        )
+
+        # Update attention
+        grounding_attn_o2i_refined = grounding_attn_o2i * refined_labels
+
+        # Create a copy to return (avoid modifying shared memory if that were an issue)
+        new_item = item.copy()
+        new_item['grounding_attn_o2i'] = grounding_attn_o2i_refined
+        return new_item
+
+    except Exception as e:
+        # Catch generic errors to prevent pool crash
+        # print(f"Error processing item: {e}")
+        return None
+
+
+# --- Main Driver ---
+def denoise_textual_attn_multiprocess(src_pkl_path, dataset_path, dst_pkl_path, num_workers=None):
+    print(f"Loading pickle from {src_pkl_path}...")
+    tmp_pkl_content = pickle.load(open(src_pkl_path, "rb"))
+
+    # Pre-filtering
+    original_len = len(tmp_pkl_content)
+    tmp_pkl_content = [item for item in tmp_pkl_content if "textvqa" not in item['image'].lower()]
+    print(f"Filtered {original_len} -> {len(tmp_pkl_content)} samples (Removed TextVQA).")
+
+    dst_pkl_content = []
+
+    # Determine workers (leave 2 cores free for system stability if possible)
+    if num_workers is None:
+        num_workers = max(1, cpu_count() - 2)
+
+    print(f"Starting multiprocessing pool with {num_workers} workers...")
+
+    # Partial function to fix the 'dataset_path' argument
+    worker_func = partial(process_single_sample, dataset_path=dataset_path)
+
+    with Pool(num_workers) as pool:
+        # imap allows us to use tqdm for a progress bar
+        # chunksize can be tuned, but default is usually fine for this data size
+        results = list(tqdm(pool.imap(worker_func, tmp_pkl_content), total=len(tmp_pkl_content)))
+
+    # Filter out failures (None returns)
+    valid_results = [res for res in results if res is not None]
+
+    print(f"Processing complete. Success: {len(valid_results)}, Failed/Skipped: {len(results) - len(valid_results)}")
+
+    print(f"Saving to {dst_pkl_path}...")
+    pickle.dump(valid_results, open(dst_pkl_path, "wb"))
+    print("Done.")
+
+
+if __name__ == "__main__":
+    # Ensure this block is used to prevent recursive process spawning on some OS
+    # denoise_textual_attn_multiprocess(
+    #     src_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_576res_247k.pkl',
+    #     dataset_path='/home/yuheng/datasets/',
+    #     dst_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_576res_229k_denoised.pkl',
+    #     num_workers=16  # Adjust based on your machine (e.g., 8, 16, 32)
+    # )
+
+    denoise_textual_attn_multiprocess(
+        src_pkl_path='/root/epfs/sproject/datasets/qwen2_5vl_pseudo_7b_576res_release.pkl',
+        dataset_path='/root/epfs/sproject/datasets/',
+        dst_pkl_path='/root/epfs/sproject/datasets/qwen2_5vl_pseudo_7b_576res_release_denoised.pkl',
+        num_workers=16  # Adjust based on your machine (e.g., 8, 16, 32)
+    )
+
+    # denoise_textual_attn_multiprocess(
+    #     src_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_576res_release.pkl',
+    #     dataset_path='//home/yuheng/datasets/',
+    #     dst_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_576res_release_denoised.pkl',
+    #     num_workers=16  # Adjust based on your machine (e.g., 8, 16, 32)
+    # )

QWENVL-PRIVATE/analysis/denoise_textual_file_wdet.py

@@ -0,0 +1,330 @@
+import sys
+import os
+import pickle
+import numpy as np
+import cv2
+import warnings
+from PIL import Image
+import torch
+import multiprocessing as mp
+from multiprocessing import Pool, Manager
+from functools import partial
+from tqdm import tqdm
+
+# Add the parent directory to sys.path to import qwen_src
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from qwen_src.mm_utils import expand2square, get_foreground_mask
+from doctr.models import ocr_predictor
+
+# --- GLOBAL LOADER ---
+_GLOBAL_PREDICTOR = None
+
+
+def get_doctr_predictor():
+    """
+    Initializes the DocTR model once per worker process on the assigned GPU.
+    """
+    global _GLOBAL_PREDICTOR
+
+    # Suppress the specific torch load warning from doctr
+    warnings.filterwarnings("ignore", category=FutureWarning, module="doctr")
+    warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
+
+    if _GLOBAL_PREDICTOR is None:
+        try:
+            # 1. Initialize the architecture (loads to CPU by default)
+            predictor = ocr_predictor(det_arch='fast_base', pretrained=True)
+
+            # 2. Force Move to GPU using the High-Level API
+            # We use 'cuda:0' because CUDA_VISIBLE_DEVICES restricts this process to seeing only 1 GPU.
+            if torch.cuda.is_available():
+                # .cuda() moves the model weights AND sets the internal flag to move input images to GPU
+                predictor.cuda()
+                predictor.det_predictor.model.eval()
+
+                # --- SANITY CHECK ---
+                # This prints the actual device of the model parameters to prove it's on GPU
+                param_device = next(predictor.det_predictor.model.parameters()).device
+                # print(f"[Worker {os.getpid()}] Model loaded. confirmed on: {param_device}")
+            else:
+                print(f"[Worker {os.getpid()}] WARNING: CUDA not available. Running on CPU.")
+
+            _GLOBAL_PREDICTOR = predictor
+
+        except Exception as e:
+            print(f"[Worker {os.getpid()}] Error loading model: {e}")
+            return None
+
+    return _GLOBAL_PREDICTOR
+
+
+def refine_pseudo_labels_with_doctr(original_image, pseudo_labels,
+                                    precomputed_boxes=None,
+                                    grid_size=(24, 24),
+                                    overlap_threshold=0.15,
+                                    predictor=None):
+    if isinstance(original_image, Image.Image):
+        img_arr = np.array(original_image)
+    else:
+        img_arr = original_image.copy()
+
+    H, W = img_arr.shape[:2]
+
+    # --- Retrieve or Compute Boxes ---
+    boxes = []
+
+    if precomputed_boxes is not None:
+        boxes = precomputed_boxes
+    else:
+        # Initialize/Get Model
+        if predictor is None:
+            doc_model = get_doctr_predictor()
+        else:
+            doc_model = predictor
+
+        # Run Inference
+        if doc_model is not None:
+            # The predictor automatically handles moving numpy/cpu inputs to the GPU
+            # if .cuda() was called during initialization.
+            result = doc_model([img_arr])
+
+            for block in result.pages[0].blocks:
+                for line in block.lines:
+                    for word in line.words:
+                        (min_x, min_y), (max_x, max_y) = word.geometry
+                        boxes.append([
+                            int(min_x * W), int(min_y * H),
+                            int(max_x * W), int(max_y * H)
+                        ])
+
+    # --- Create Text Mask ---
+    text_mask = np.zeros((H, W), dtype=np.uint8)
+    for box in boxes:
+        cv2.rectangle(text_mask, (box[0], box[1]), (box[2], box[3]), 1, thickness=-1)
+
+    # -- -Refine Labels ---
+    h_grid, w_grid = grid_size
+    if pseudo_labels.shape != grid_size:
+        pseudo_labels = pseudo_labels.reshape(h_grid, w_grid)
+
+    refined_labels = pseudo_labels.copy()
+    step_h, step_w = H // h_grid, W // w_grid
+    rows, cols = np.where(pseudo_labels == 1)
+
+    for r, c in zip(rows, cols):
+        y_start, y_end = r * step_h, min((r + 1) * step_h, H)
+        x_start, x_end = c * step_w, min((c + 1) * step_w, W)
+
+        mask_patch = text_mask[y_start:y_end, x_start:x_end]
+        if mask_patch.size > 0:
+            ratio = np.sum(mask_patch) / mask_patch.size
+            if ratio < overlap_threshold:
+                refined_labels[r, c] = 0
+        else:
+            refined_labels[r, c] = 0
+
+    return refined_labels, boxes
+
+
+# --- Worker Functions ---
+
+def gpu_worker_init(queue):
+    """
+    Assigns a unique GPU to the worker process.
+    """
+    try:
+        gpu_id = queue.get()
+        # Mask the environment so this process only sees this specific GPU
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+
+        # Suppress warnings inside the worker process as well
+        warnings.filterwarnings("ignore", category=FutureWarning)
+    except Exception as e:
+        print(f"Worker init failed: {e}")
+
+
+def process_single_sample_packed(args):
+    """
+    Unpacks arguments and processes the sample.
+    args: (item, dataset_path, precomputed_boxes)
+    """
+    item, dataset_path, precomputed_boxes = args
+
+    try:
+        img_path = item['image']
+
+        # --- GQA / NATURAL IMAGE HANDLING ---
+        # Keep GQA samples intact, but return None for boxes (no cache update)
+        if 'gqa' in img_path.lower():
+            return item, img_path, None
+
+        # --- PROCESSING FOR TEXT IMAGES ---
+        full_img_path = os.path.join(dataset_path, img_path) if dataset_path else img_path
+        grounding_attn_o2i = item['grounding_attn_o2i']
+
+        try:
+            image = Image.open(full_img_path).convert('RGB')
+        except Exception:
+            return None, None, None
+
+        # Prepare Mask
+        known_fg_mask = get_foreground_mask(image.size, grounding_attn_o2i.shape[0], min_dimension_size=1)
+        tmp_attn = grounding_attn_o2i.copy()
+        tmp_attn[~known_fg_mask] = 0.
+        tmp_attn_binary = (tmp_attn > tmp_attn.max() * 0.05).astype(np.float32)
+
+        # Denoise
+        square_image = expand2square(image, background_color=(127, 127, 127))
+
+        refined_labels, final_boxes = refine_pseudo_labels_with_doctr(
+            original_image=square_image,
+            pseudo_labels=tmp_attn_binary,
+            precomputed_boxes=precomputed_boxes,
+            grid_size=(grounding_attn_o2i.shape[0], grounding_attn_o2i.shape[1]),
+            overlap_threshold=0.15
+        )
+
+        new_item = item.copy()
+        new_item['grounding_attn_o2i'] = grounding_attn_o2i * refined_labels
+
+        return new_item, img_path, final_boxes
+
+    except Exception as e:
+        # print(f"Error in worker: {e}")
+        return None, None, None
+
+
+def denoise_textual_attn_multiprocess(
+    src_pkl_path,
+    dataset_path,
+    dst_pkl_path,
+    ocr_cache_file,
+    num_gpus=None,
+    filter_substr="docvqa",
+):
+    # 1. Setup GPUs
+    if num_gpus is None:
+        num_gpus = torch.cuda.device_count()
+        if num_gpus == 0: num_gpus = 1
+
+    print(f"Using {num_gpus} workers (GPUs).")
+
+    # 2. Load Dataset
+    print(f"Loading dataset from {src_pkl_path}...")
+    tmp_pkl_content = pickle.load(open(src_pkl_path, "rb"))
+
+    # Optional: Filter explicit datasets if needed
+    original_len = len(tmp_pkl_content)
+    # # Removing TextVQA as per previous requests, keeping GQA handled in worker
+    # tmp_pkl_content = [item for item in tmp_pkl_content if "textvqa" not in item['image'].lower()]
+    if filter_substr:
+        tmp_pkl_content = [item for item in tmp_pkl_content if filter_substr in item['image'].lower()]
+
+    print(f"Filtered {original_len} -> {len(tmp_pkl_content)} samples.")
+
+    # 3. Load/Init OCR Cache (Single File)
+    ocr_cache = {}
+    if ocr_cache_file and os.path.exists(ocr_cache_file):
+        print(f"Loading existing OCR cache from {ocr_cache_file}...")
+        try:
+            with open(ocr_cache_file, 'rb') as f:
+                ocr_cache = pickle.load(f)
+            print(f"Loaded {len(ocr_cache)} cached entries.")
+        except Exception as e:
+            print(f"Failed to load cache: {e}. Starting fresh.")
+
+    # 4. Prepare Tasks
+    tasks = []
+    for item in tmp_pkl_content:
+        img_path = item['image']
+        cached_boxes = ocr_cache.get(img_path, None)
+        tasks.append((item, dataset_path, cached_boxes))
+
+    # 5. Start Processing
+    m = Manager()
+    gpu_queue = m.Queue()
+    for i in range(num_gpus): gpu_queue.put(i)
+
+    print("Starting pool...")
+    valid_results = []
+    cache_updated = False
+
+    # --- FORCE SPAWN ---
+    # This prevents the parent process's state from corrupting the worker's CUDA context
+    try:
+        mp.set_start_method('spawn', force=True)
+    except RuntimeError:
+        pass
+
+    with Pool(processes=num_gpus, initializer=gpu_worker_init, initargs=(gpu_queue,)) as pool:
+        for result_item, img_path, boxes in tqdm(pool.imap(process_single_sample_packed, tasks), total=len(tasks)):
+
+            if result_item is not None:
+                valid_results.append(result_item)
+
+                # Only update cache if we actually have boxes (i.e. not GQA)
+                if img_path and boxes is not None:
+                    if img_path not in ocr_cache:
+                        ocr_cache[img_path] = boxes
+                        cache_updated = True
+
+    print(f"Processing complete. Valid items: {len(valid_results)}")
+
+    # 6. Save Outputs
+    print(f"Saving denoised dataset to {dst_pkl_path}...")
+    with open(dst_pkl_path, "wb") as f:
+        pickle.dump(valid_results, f)
+
+    if ocr_cache_file:
+        if cache_updated or not os.path.exists(ocr_cache_file):
+            print(f"Saving updated OCR cache ({len(ocr_cache)} entries) to {ocr_cache_file}...")
+            with open(ocr_cache_file, "wb") as f:
+                pickle.dump(ocr_cache, f)
+        else:
+            print("OCR Cache was not modified, skipping save.")
+
+    print("Done.")
+
+if __name__ == "__main__":
+    # --- CRITICAL FIX FOR GPU MULTIPROCESSING ---
+    # This must be the very first thing that happens in __main__
+    # 'spawn' creates a fresh process, ensuring CUDA is not corrupted by the parent.
+    try:
+        mp.set_start_method('spawn', force=True)
+    except RuntimeError as e:
+        print(f"Start method setting warning: {e}")
+
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Denoise textual-image pseudo labels with DocTR detection."
+    )
+    parser.add_argument("--src-pkl-path", type=str, required=True, help="Source pkl file path.")
+    parser.add_argument("--dataset-path", type=str, default=None, help="Root dataset path for image loading.")
+    parser.add_argument("--dst-pkl-path", type=str, required=True, help="Output pkl file path.")
+    parser.add_argument("--ocr-cache-file", type=str, default=None, help="Path to OCR cache pkl.")
+    parser.add_argument("--num-gpus", type=int, default=None, help="Number of GPUs to use.")
+    parser.add_argument(
+        "--filter-substr",
+        type=str,
+        default="docvqa",
+        help="Only keep samples whose image path contains this substring (case-insensitive).",
+    )
+    args = parser.parse_args()
+
+    denoise_textual_attn_multiprocess(
+        src_pkl_path=args.src_pkl_path,
+        dataset_path=args.dataset_path,
+        dst_pkl_path=args.dst_pkl_path,
+        ocr_cache_file=args.ocr_cache_file,
+        num_gpus=args.num_gpus,
+        filter_substr=args.filter_substr.lower() if args.filter_substr else "",
+    )
+    # denoise_textual_attn_multiprocess(
+    #     src_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_576res_release.pkl',
+    #     dataset_path='/home/yuheng/datasets/',
+    #     dst_pkl_path='/home/yuheng/datasets/qwen2_5vl_pseudo_7b_576res_release_denoised.pkl',
+    #     ocr_cache_file='/home/yuheng/datasets/ocr_boxes_cache_full.pkl',
+    #     num_gpus=4
+    # )

QWENVL-PRIVATE/analysis/failure_case_analysis.py

@@ -0,0 +1,302 @@
+import json
+import os
+
+
+def load_jsonl(path):
+    """
+    Loads a JSONL file into a dictionary keyed by 'doc_id'.
+    """
+    data = {}
+    try:
+        with open(path, 'r', encoding='utf-8') as f:
+            for i, line in enumerate(f):
+                try:
+                    record = json.loads(line)
+                    # Use 'doc_id' as the unique key for comparison
+                    if 'doc_id' not in record:
+                        print(f"Warning: Missing 'doc_id' in {path} at line {i + 1}. Skipping.")
+                        continue
+                    data[record['doc_id']] = record
+                except json.JSONDecodeError:
+                    print(f"Error: Could not parse JSON in {path} at line {i + 1}. Skipping.")
+    except FileNotFoundError:
+        print(f"Error: File not found at {path}")
+        return None
+    except Exception as e:
+        print(f"Error opening file {path}: {e}")
+        return None
+
+    print(f"Successfully loaded {len(data)} records from {path}.")
+    return data
+
+
+def find_failure_cases_textvqa(ours_output_path, baseline_output_path, score_key="exact_match"):
+    """
+    Compares two output files and finds items where 'ours' is wrong (0.0)
+    and 'baseline' is right (1.0).
+
+    Args:
+        ours_output_path (str): Path to your model's output JSONL.
+        baseline_output_path (str): Path to the baseline model's output JSONL.
+        score_key (str): The key in the JSON to find the score dictionary.
+                         Default is "vstar_overall_acc".
+
+    Returns:
+        list: A list of dictionaries, each containing details of a failure case.
+    """
+    failure_cases = []
+
+    # Load the baseline results into a dictionary for fast lookup
+    baseline_results = load_jsonl(baseline_output_path)
+    if baseline_results is None:
+        return []
+
+    print(f"\nProcessing {ours_output_path} against baseline...")
+
+    try:
+        with open(ours_output_path, 'r', encoding='utf-8') as f:
+            for i, line in enumerate(f):
+                try:
+                    our_record = json.loads(line)
+                    if 'doc_id' not in our_record:
+                        print(f"Warning: Missing 'doc_id' in 'ours' file at line {i + 1}. Skipping.")
+                        continue
+                    doc_id = our_record['doc_id']
+
+                    # Check if this doc_id exists in the baseline
+                    if doc_id not in baseline_results:
+                        print(f"Warning: doc_id {doc_id} from 'ours' file not found in baseline. Skipping.")
+                        continue
+                    baseline_record = baseline_results[doc_id]
+                    # Check if score keys exist before accessing
+                    if score_key not in our_record :
+                        print(f"Warning: Missing score key '{score_key}' in 'ours' file for doc_id {doc_id}. Skipping.")
+                        continue
+                    if score_key not in baseline_record:
+                        print(
+                            f"Warning: Missing score key '{score_key}' in baseline file for doc_id {doc_id}. Skipping.")
+                        continue
+                    # Get the scores
+                    if type(our_record[score_key]) is not dict or type(baseline_record[score_key]) is not dict:
+                        our_score = our_record[score_key]
+                        baseline_score = baseline_record[score_key]
+                    else:
+                        our_score = our_record[score_key]['score']
+                        baseline_score = baseline_record[score_key]['score']
+
+                    # The core logic: Find where we failed and baseline succeeded
+                    if our_score < 0.3 and baseline_score > 0.7:
+                        # Format a clean dictionary for case study
+                        failure_info = {
+                            "doc_id": doc_id,
+                            "question": our_record['input'],
+                            "ground_truth": baseline_record['filtered_resps'][0],
+                            "our_prediction": our_record['filtered_resps'][0],#[score_key].get('prediction', 'N/A'),
+                            "baseline_prediction": baseline_record['filtered_resps'][0],
+                            "full_our_record": our_record  # Optional: include for full context
+                        }
+                        failure_cases.append(failure_info)
+
+                except json.JSONDecodeError:
+                    print(f"Error: Could not parse JSON in 'ours' file at line {i + 1}. Skipping.")
+                except KeyError:
+                    # This catches if a sub-key like 'score' is missing
+                    print(f"Error: Missing expected key structure in 'ours' file at line {i + 1}. Skipping.")
+
+    except FileNotFoundError:
+        print(f"Error: File not found at {ours_output_path}")
+        return []
+    except Exception as e:
+        print(f"Error opening file {ours_output_path}: {e}")
+        return []
+
+    print(f"\n--- Analysis Complete ---")
+    print(f"Found {len(failure_cases)} failure cases where 'ours' was wrong and 'baseline' was right.")
+    return failure_cases
+
+
+def find_failure_cases_vstar(ours_output_path, baseline_output_path, score_key="vstar_overall_acc"):
+    """
+    Compares two output files and finds items where 'ours' is wrong (0.0)
+    and 'baseline' is right (1.0).
+
+    Args:
+        ours_output_path (str): Path to your model's output JSONL.
+        baseline_output_path (str): Path to the baseline model's output JSONL.
+        score_key (str): The key in the JSON to find the score dictionary.
+                         Default is "vstar_overall_acc".
+
+    Returns:
+        list: A list of dictionaries, each containing details of a failure case.
+    """
+    failure_cases = []
+
+    # Load the baseline results into a dictionary for fast lookup
+    baseline_results = load_jsonl(baseline_output_path)
+    if baseline_results is None:
+        return []
+
+    print(f"\nProcessing {ours_output_path} against baseline...")
+
+    try:
+        with open(ours_output_path, 'r', encoding='utf-8') as f:
+            for i, line in enumerate(f):
+                try:
+                    our_record = json.loads(line)
+
+                    if 'doc_id' not in our_record:
+                        print(f"Warning: Missing 'doc_id' in 'ours' file at line {i + 1}. Skipping.")
+                        continue
+
+                    doc_id = our_record['doc_id']
+
+                    # Check if this doc_id exists in the baseline
+                    if doc_id not in baseline_results:
+                        print(f"Warning: doc_id {doc_id} from 'ours' file not found in baseline. Skipping.")
+                        continue
+
+                    baseline_record = baseline_results[doc_id]
+
+                    # Check if score keys exist before accessing
+                    if score_key not in our_record or 'score' not in our_record[score_key]:
+                        print(f"Warning: Missing score key '{score_key}' in 'ours' file for doc_id {doc_id}. Skipping.")
+                        continue
+
+                    if score_key not in baseline_record or 'score' not in baseline_record[score_key]:
+                        print(
+                            f"Warning: Missing score key '{score_key}' in baseline file for doc_id {doc_id}. Skipping.")
+                        continue
+
+                    # Get the scores
+                    our_score = our_record[score_key]['score']
+                    baseline_score = baseline_record[score_key]['score']
+
+                    # The core logic: Find where we failed and baseline succeeded
+                    if our_score == 0.0 and baseline_score == 1.0:
+                        # Format a clean dictionary for case study
+                        failure_info = {
+                            "doc_id": doc_id,
+                            "question": our_record.get('doc', {}).get('text', 'N/A'),
+                            "ground_truth": our_record.get('target', 'N/A'),
+                            "our_prediction": our_record[score_key].get('prediction', 'N/A'),
+                            "baseline_prediction": baseline_record[score_key].get('prediction', 'N/A'),
+                            "full_our_record": our_record  # Optional: include for full context
+                        }
+                        failure_cases.append(failure_info)
+
+                except json.JSONDecodeError:
+                    print(f"Error: Could not parse JSON in 'ours' file at line {i + 1}. Skipping.")
+                except KeyError:
+                    # This catches if a sub-key like 'score' is missing
+                    print(f"Error: Missing expected key structure in 'ours' file at line {i + 1}. Skipping.")
+
+    except FileNotFoundError:
+        print(f"Error: File not found at {ours_output_path}")
+        return []
+    except Exception as e:
+        print(f"Error opening file {ours_output_path}: {e}")
+        return []
+
+    print(f"\n--- Analysis Complete ---")
+    print(f"Found {len(failure_cases)} failure cases where 'ours' was wrong and 'baseline' was right.")
+    return failure_cases
+
+# --- Example Usage ---
+if __name__ == "__main__":
+
+    # --- 1. Create dummy data files for the example ---
+
+    # This example has 3 cases:
+    # 1. We fail (B), baseline succeeds (A). -> This should be found.
+    # 2. We succeed (A), baseline fails (C). -> This should be ignored.
+    # 3. Both succeed (A). -> This should be ignored.
+
+
+
+    # # Define file paths
+    # baseline_file = "./logs/liuhaotian__llava-v1.5-7b/20250902_190647_samples_textvqa_val.jsonl"
+    # ours_file = "./logs/checkpoints__llava-v1.5-7b-roi-K15T3-152k-v1bf16Mheads-bg0p05fg0p2-filled/20251113_113923_samples_textvqa_val.jsonl"
+    #
+    # #failures = find_failure_cases_textvqa(ours_file, baseline_file, 'exact_match')
+    # failures = find_failure_cases_textvqa(baseline_file, ours_file , 'exact_match')
+
+
+    #baseline_file = "./logs/liuhaotian__llava-v1.5-7b/20250902_201850_samples_vstar_bench.jsonl"
+    #ours_file = "./logs/checkpoints__llava-v1.5-7b-roi-K15T3-152k-v1bf16Mheads-bg0p05fg0p2-filled/20251113_130922_samples_vstar_bench.jsonl"
+
+    # baseline_file = "../Qwen2.5-VL/logs/Qwen__Qwen2.5-VL-7B-Instruct/20251113_135313_samples_vstar_bench.jsonl"
+    # ours_file = "../Qwen2.5-VL/logs/output__qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled/20251113_134244_samples_vstar_bench.jsonl"
+    #
+    # failures = find_failure_cases_vstar(baseline_file, ours_file)
+    # failure_ids = [case['doc_id'] for case in failures]
+
+    ###### TextVQA Failure Case Analysis
+    #baseline_file_textvqa = "logs/output__qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled/20251209_181751_samples_textvqa_val.jsonl"
+    baseline_file_textvqa = "logs/Qwen__Qwen3-VL-4B-Instruct/20260122_115842_samples_textvqa_val.jsonl"
+
+    ours_file_textvqa = "logs/output__qwen3vl-4b-roi-K24T3-185k-v1bf16-twiginit/20260122_120855_samples_textvqa_val.jsonl"
+    failures_baseline_textvqa = find_failure_cases_textvqa(baseline_file_textvqa, ours_file_textvqa, 'exact_match')
+    failure_ids_baseline_textvqa = [case['doc_id'] for case in failures_baseline_textvqa]
+
+    failures_ours_textvqa = find_failure_cases_textvqa(ours_file_textvqa, baseline_file_textvqa, 'exact_match')
+    failure_ids_ours_textvqa = [case['doc_id'] for case in failures_ours_textvqa]
+    failure_ids_ours_textvqa_dict = {case['doc_id']: case for case in failures_ours_textvqa}
+    ### print some results
+    for i, case in enumerate(failures_ours_textvqa[:10]):
+        print(f"\nCase {i + 1} (doc_id: {case['doc_id']})")
+        print(f"  Question: {case['question']}")
+        print(f"  Ground Truth: {case['ground_truth']}")
+        print(f"  Our (Wrong) Prediction: {case['our_prediction']}")
+        print(f"  Baseline (Correct) Prediction: {case['baseline_prediction']}")
+
+
+
+    # ###### InfoVQA Failure Case Analysis
+    # baseline_file_infovqa = "../Qwen2.5-VL/logs/Qwen__Qwen2.5-VL-7B-Instruct/20251204_093838_samples_infovqa_val.jsonl"
+    # ours_file_infovqa = "../Qwen2.5-VL/logs/output__qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled/20251215_131620_samples_infovqa_val.jsonl"
+    #
+    # failures_baseline_infovqa = find_failure_cases_textvqa(baseline_file_infovqa, ours_file_infovqa, 'anls')
+    # failure_ids_baseline_infovqa = [case['doc_id'] for case in failures_baseline_infovqa]
+    #
+    # failures_ours_infovqa = find_failure_cases_textvqa(ours_file_infovqa, baseline_file_infovqa, 'anls')
+    # failure_ids_ours_infovqa = [case['doc_id'] for case in failures_ours_infovqa]
+    #
+    # ##### DocVQA Failure Case Analysis
+    # baseline_file_docvqa = "../Qwen2.5-VL/logs/Qwen__Qwen2.5-VL-7B-Instruct/20251204_095001_samples_docvqa_val.jsonl"
+    # ours_file_docvqa = "../Qwen2.5-VL/logs/output__qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled/20251210_143202_samples_docvqa_val.jsonl"
+    #
+    # failures_baseline_docvqa = find_failure_cases_textvqa(baseline_file_docvqa, ours_file_docvqa, 'anls')
+    # failure_ids_baseline_docvqa = [case['doc_id'] for case in failures_baseline_docvqa]
+    # failures_ours_docvqa = find_failure_cases_textvqa(ours_file_docvqa, baseline_file_docvqa, 'anls')
+    # failure_ids_ours_docvqa = [case['doc_id'] for case in failures_ours_docvqa]
+    #
+    # ##### ChartQA Failure Case Analysis
+    # baseline_file_chartqa = "../Qwen2.5-VL/logs/Qwen__Qwen2.5-VL-7B-Instruct/20251204_094359_samples_chartqa.jsonl"
+    # ours_file_chartqa = "../Qwen2.5-VL/logs/output__qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled/20251210_142453_samples_chartqa.jsonl"
+    #
+    # failures_baseline_chartqa = find_failure_cases_textvqa(baseline_file_chartqa, ours_file_chartqa, 'relaxed_overall')
+    # failure_ids_baseline_chartqa = [case['doc_id'] for case in failures_baseline_chartqa]
+    # failures_ours_chartqa = find_failure_cases_textvqa(ours_file_chartqa, baseline_file_chartqa, 'relaxed_overall')
+    # failure_ids_ours_chartqa = [case['doc_id'] for case in failures_ours_chartqa]
+    #
+    #
+    # ###### OCRBench Failure Case Analysis
+    # baseline_file_ocrbench = "../Qwen2.5-VL/logs/Qwen__Qwen2.5-VL-7B-Instruct/20251204_094800_samples_ocrbench.jsonl"
+    # ours_file_ocrbench = "../Qwen2.5-VL/logs/output__qwen2_5vl-7b-roi-K18T3-134k-gatedv2PB-wGTwRoI-twiginit-filled/20251210_142936_samples_ocrbench.jsonl"
+    # failures_baseline_ocrbench = find_failure_cases_textvqa(baseline_file_ocrbench, ours_file_ocrbench, 'ocrbench_accuracy')
+    # failure_ids_baseline_ocrbench = [case['doc_id'] for case in failures_baseline_ocrbench]
+    # failures_ours_ocrbench = find_failure_cases_textvqa(ours_file_ocrbench, baseline_file_ocrbench, 'ocrbench_accuracy')
+    # failure_ids_ours_ocrbench = [case['doc_id'] for case in failures_ours_ocrbench]
+    #
+
+    pass
+
+
+    # if failure_ids_ours:
+    #     print("\n--- Failure Cases ---")
+    #     for i, case in enumerate(failure_ids_ours):
+    #         print(f"\nCase {i + 1} (doc_id: {case['doc_id']})")
+    #         print(f"  Question: {case['question']}")
+    #         print(f"  Ground Truth: {case['ground_truth']}")
+    #         print(f"  Our (Wrong) Prediction: {case['our_prediction']}")
+    #         print(f"  Baseline (Correct) Prediction: {case['baseline_prediction']}")

QWENVL-PRIVATE/analysis/filter_roi_confusion.py

@@ -0,0 +1,92 @@
+import json
+import argparse
+import os
+
+
+def filter_improvements(input_path, output_path, exclude_datasets):
+    # 1. Define keywords that indicate the model is confused about image count
+    EXCLUDE_KEYWORDS = [
+        "images",
+        "second image",
+        "image 1",
+        "image 2",
+        "two images"
+    ]
+
+    print(f"Loading data from: {input_path}")
+    print(f"Excluding datasets: {exclude_datasets}")
+
+    if not os.path.exists(input_path):
+        print(f"Error: File not found at {input_path}")
+        return
+
+    with open(input_path, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+
+    original_count = len(data)
+    filtered_data = []
+
+    removed_keyword_count = 0
+    removed_dataset_count = 0
+
+    print(f"Processing {original_count} entries...")
+
+    for entry in data:
+        roi_resp = entry.get('roi_response', '').lower()
+        dataset_name = entry.get('dataset', '')
+
+        # Check 1: Is this dataset in the exclusion list?
+        if dataset_name in exclude_datasets:
+            removed_dataset_count += 1
+            continue
+
+        # Check 2: Does the response contain confusion keywords?
+        should_exclude_keyword = any(keyword in roi_resp for keyword in EXCLUDE_KEYWORDS)
+        if should_exclude_keyword:
+            removed_keyword_count += 1
+            continue
+
+        # Keep entry if it passes both checks
+        filtered_data.append(entry)
+
+    # 2. Save the cleaned data
+    with open(output_path, 'w', encoding='utf-8') as f:
+        json.dump(filtered_data, f, indent=2)
+
+    # 3. Print Statistics
+    total_removed = removed_keyword_count + removed_dataset_count
+    print("-" * 30)
+    print(f"Original entries:        {original_count}")
+    print(f"Removed (Keywords):      {removed_keyword_count}")
+    print(f"Removed (Dataset match): {removed_dataset_count}")
+    print(f"Total Removed:           {total_removed}")
+    print(f"Remaining entries:       {len(filtered_data)}")
+    print("-" * 30)
+    print(f"Filtered file saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Filter ROI confusion and specific datasets from improvements.json")
+    parser.add_argument('--input', type=str, required=True,
+                        help="Path to the original improvements JSON file")
+    parser.add_argument('--output', type=str, required=True,
+                        help="Path to save the filtered JSON file")
+
+    # New argument for dataset exclusion
+    parser.add_argument('--exclude-datasets', type=str, default="vstar_coco_spatial",
+                        help="Comma-separated list of datasets to remove (e.g. 'vstar_coco_spatial,other_dataset')")
+
+    args = parser.parse_args()
+
+    # Convert comma-separated string to a list of strings
+    exclude_list = [d.strip() for d in args.exclude_datasets.split(',') if d.strip()]
+
+    filter_improvements(args.input, args.output, exclude_list)
+
+
+'''
+python analysis/filter_roi_confusion.py \
+    --input /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-32b-judge.json \
+    --output /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-clean-32b-judge-wo-spatial.json \
+    --exclude-datasets "vstar_coco_spatial"
+'''

QWENVL-PRIVATE/analysis/filter_roi_fails.py

@@ -0,0 +1,438 @@
+import os
+import json
+import pickle
+import argparse
+import math
+import torch
+import multiprocessing as mp
+import re
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoProcessor
+from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
+try:
+    from qwen_src.qwen3_vl import Qwen3VLForConditionalGeneration
+except Exception:
+    Qwen3VLForConditionalGeneration = None
+from tqdm import tqdm
+import glob
+
+
+# --- 1. Dataset & Collator ---
+
+class JudgeDataset(Dataset):
+    def __init__(self, items):
+        self.items = items
+
+    def __len__(self):
+        return len(self.items)
+
+    def __getitem__(self, idx):
+        item = self.items[idx]
+
+        # 1. Format Ground Truth
+        gt = item['ground_truth']
+        if isinstance(gt, list):
+            gt_str = ", ".join([str(x) for x in gt])
+        else:
+            gt_str = str(gt)
+
+        # 2. Judge Prompt
+        prompt_text = (
+            f"You are an intelligent evaluator for a Visual Question Answering task.\n"
+            f"Compare the Model Predictions against the Ground Truth.\n\n"
+
+            f"--- DATA ---\n"
+            f"Question: {item['question']}\n"
+            f"Ground Truth: {gt_str}\n"
+            f"Model A Prediction: {item['base_resp']}\n"
+            f"Model B Prediction: {item['roi_resp']}\n\n"
+
+            f"--- CRITERIA ---\n"
+            f"1. **Containment**: If the Ground Truth is a short value (e.g., '6000') and the Model Prediction contains it correctly within a sentence (e.g., 'The value is 6000'), mark it as **Yes**.\n"
+            f"2. **Synonyms**: Accept standard synonyms (e.g., 'Yes' = 'True', '10%' = '0.1').\n"
+            f"3. **Formatting**: Ignore punctuation or capitalization differences (e.g., 'april' = 'April').\n"
+            f"4. **Contradiction**: Only mark 'No' if the prediction is factually wrong or missing the key information.\n\n"
+
+            f"--- TASK ---\n"
+            f"Does Model A provide the correct answer? (Yes/No)\n"
+            f"Does Model B provide the correct answer? (Yes/No)\n\n"
+
+            f"Respond in this exact format:\n"
+            f"A: [Yes/No]\n"
+            f"B: [Yes/No]"
+        )
+
+        messages = [
+            {"role": "system",
+             "content": "You are a fair and robust evaluator. You focus on semantic correctness over strict string matching."},
+            {"role": "user", "content": prompt_text}
+        ]
+
+        return {
+            "messages": messages,
+            "original_item": item
+        }
+
+
+def infer_model_type(model_path, model_type):
+    if model_type and model_type != "auto":
+        return model_type
+    lower = model_path.lower()
+    if "qwen3" in lower:
+        return "qwen3"
+    if "qwen2.5" in lower or "qwen2_5" in lower or "qwen2-5" in lower:
+        return "qwen2_5"
+    return "qwen2_5"
+
+
+def resolve_processor_path(model_path, model_type, processor_path=None):
+    if processor_path:
+        return processor_path
+    lower = model_path.lower()
+    public_prefix = "/root/public/Qwen"
+
+    def to_repo_path(name):
+        if "/root/public" in model_path:
+            return f"{public_prefix}/{name}"
+        return f"Qwen/{name}"
+
+    if model_type == "qwen3":
+        if "8b" in lower:
+            return to_repo_path("Qwen3-VL-8B-Instruct")
+        if "4b" in lower:
+            return to_repo_path("Qwen3-VL-4B-Instruct")
+    if "7b" in lower:
+        return to_repo_path("Qwen2.5-VL-7B-Instruct")
+    return model_path
+
+
+def resolve_model_class(model_type):
+    if model_type == "qwen3":
+        if Qwen3VLForConditionalGeneration is None:
+            raise RuntimeError("Qwen3VLForConditionalGeneration not available.")
+        return Qwen3VLForConditionalGeneration
+    return Qwen2_5_VLForConditionalGeneration
+
+
+class JudgeCollator:
+    def __init__(self, processor_path):
+        self.processor_path = processor_path
+        self.processor = None
+
+    def __call__(self, batch):
+        if self.processor is None:
+            self.processor = AutoProcessor.from_pretrained(
+                self.processor_path,
+                trust_remote_code=True,
+                padding_side='left',
+                use_fast=True
+            )
+
+        raw_messages = [item['messages'] for item in batch]
+        original_items = [item['original_item'] for item in batch]
+
+        text_inputs = [
+            self.processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+            for msg in raw_messages
+        ]
+
+        inputs = self.processor(
+            text=text_inputs,
+            padding=True,
+            return_tensors="pt"
+        )
+
+        return {
+            "inputs": inputs,
+            "items": original_items
+        }
+
+
+# --- 2. Worker Functions ---
+
+def worker_process(chunk, gpu_id, args):
+    device = f'cuda:{gpu_id}'
+
+    # Storage for results
+    regressions_found = []  # Base: Right, ROI: Wrong
+    improvements_found = []  # Base: Wrong, ROI: Right
+
+    # Setup Dataset & Loader
+    dataset = JudgeDataset(chunk)
+    collator = JudgeCollator(args.processor_path)
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        collate_fn=collator,
+        pin_memory=True
+    )
+
+    if len(dataset) == 0:
+        return {"regressions": [], "improvements": []}
+
+    try:
+        # Load Model
+        print(f"Worker {gpu_id}: Loading Judge Model...")
+        model_class = resolve_model_class(args.model_type)
+        model = model_class.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device
+        )
+        model.eval()
+
+        tokenizer = collator.processor.tokenizer if collator.processor else None
+        if tokenizer is None:
+            tokenizer = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True).tokenizer
+
+        print(f"Worker {gpu_id}: Starting Judging...")
+
+        for batch_data in tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id):
+            inputs = batch_data['inputs']
+            items = batch_data['items']
+
+            inputs = inputs.to(device)
+
+            with torch.no_grad():
+                output_ids = model.generate(
+                    **inputs,
+                    max_new_tokens=32,
+                    do_sample=False,
+                    pad_token_id=tokenizer.eos_token_id
+                )
+
+            input_len = inputs.input_ids.shape[1]
+            generated_ids = output_ids[:, input_len:]
+            responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+
+            for item, resp in zip(items, responses):
+                clean_resp = resp.replace("\n", " ").strip()
+
+                # Parse A (Baseline)
+                base_match = re.search(r"A\s*[:\-\.]\s*(Yes|True|Correct)", clean_resp, re.IGNORECASE)
+                base_correct = bool(base_match)
+
+                # Parse B (ROI)
+                roi_match = re.search(r"B\s*[:\-\.]\s*(Yes|True|Correct)", clean_resp, re.IGNORECASE)
+                roi_correct = bool(roi_match)
+
+                # Helper to create result dict
+                def create_result_entry():
+                    s_bbox = item.get('surrounding_bbox')
+                    if hasattr(s_bbox, 'tolist'): s_bbox = s_bbox.tolist()
+
+                    r_mask = item.get('roi_mask')
+                    mask_shape = r_mask.shape if hasattr(r_mask, 'shape') else "N/A"
+
+                    return {
+                        "dataset": item.get('dataset', 'unknown'),
+                        "line_number": item['line_number'],
+                        "question_id": item['id'],
+                        "image_id": item['image'],
+                        "question": item['question'],
+                        "ground_truth": item['ground_truth'],
+                        "baseline_response": item['base_resp'],
+                        "roi_response": item['roi_resp'],
+                        "judge_raw_output": resp,
+                        "surrounding_bbox": s_bbox,
+                        "roi_mask_shape": mask_shape
+                    }
+
+                # CASE 1: REGRESSION (Base Right, ROI Wrong)
+                if base_correct and not roi_correct:
+                    regressions_found.append(create_result_entry())
+
+                # CASE 2: IMPROVEMENT (Base Wrong, ROI Right)
+                elif not base_correct and roi_correct:
+                    improvements_found.append(create_result_entry())
+
+            del inputs, output_ids, generated_ids
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: Critical Error: {e}")
+        import traceback
+        traceback.print_exc()
+
+    return {"regressions": regressions_found, "improvements": improvements_found}
+
+
+def worker_wrapper(chunk, gpu_id, args, return_dict):
+    result = worker_process(chunk, gpu_id, args)
+    return_dict[gpu_id] = result
+
+
+# --- 3. Main Execution ---
+
+def load_all_paired_data(search_dir, file_pattern="*_qwen2.5vl-7B.pkl"):
+    all_base_data = {}
+    all_roi_data = {}
+
+    search_pattern = os.path.join(search_dir, file_pattern)
+    base_files = glob.glob(search_pattern)
+
+    print(f"Found {len(base_files)} baseline files in '{search_dir}'. Loading pairs...")
+
+    for base_file in base_files:
+        filename = os.path.basename(base_file)
+        suffix_name = file_pattern.replace("*", "")
+        dataset_name = filename.replace(suffix_name, "")
+        roi_file = base_file.replace(".pkl", "-roi.pkl")
+
+        if not os.path.exists(roi_file):
+            print(f"⚠️  Skipping {dataset_name}: ROI file not found ({roi_file})")
+            continue
+
+        try:
+            with open(base_file, 'rb') as f:
+                cur_base = pickle.load(f)
+            with open(roi_file, 'rb') as f:
+                cur_roi = pickle.load(f)
+
+            for k in cur_base.keys():
+                if k in cur_roi:
+                    unique_key = f"{dataset_name}_{k}"
+                    all_base_data[unique_key] = cur_base[k]
+                    all_roi_data[unique_key] = cur_roi[k]
+                    all_base_data[unique_key]['dataset'] = dataset_name
+                    all_roi_data[unique_key]['dataset'] = dataset_name
+
+            print(f"  + Loaded {dataset_name}: {len(cur_base)} samples")
+
+        except Exception as e:
+            print(f"❌ Error loading {dataset_name}: {e}")
+
+    return all_base_data, all_roi_data
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model-path', type=str, default='Qwen/Qwen2.5-VL-7B-Instruct')
+    parser.add_argument('--model-type', type=str, default="auto",
+                        choices=["auto", "qwen2_5", "qwen3"],
+                        help="Model family for loading (auto detects from model path).")
+    parser.add_argument('--processor-path', type=str, default="",
+                        help="Override processor path if auto detection is incorrect.")
+    parser.add_argument('--data-dir', type=str, default='./', help="Directory containing .pkl files")
+    parser.add_argument('--save-path', type=str, default="regressions.json",
+                        help="File for Regressions (Base Right, ROI Wrong)")
+    parser.add_argument('--save-path-improvements', type=str, default="improvements.json",
+                        help="File for Improvements (Base Wrong, ROI Right)")
+    parser.add_argument('--gpu-ids', type=str, default='0,1,2,3')
+    parser.add_argument('--batch-size', type=int, default=8)
+    parser.add_argument('--num_workers', type=int, default=4)
+    parser.add_argument('--file-pattern', type=str, default="*_qwen2.5vl-7B.pkl",)
+    args = parser.parse_args()
+    args.model_type = infer_model_type(args.model_path, args.model_type)
+    args.processor_path = resolve_processor_path(
+        args.model_path, args.model_type, args.processor_path
+    )
+
+    # 1. Load All Data Pairs
+    base_data, roi_data = load_all_paired_data(args.data_dir, file_pattern=args.file_pattern)
+
+    if not base_data:
+        print("No paired data found. Exiting.")
+        return
+
+    # 2. Prepare Items
+    judge_items = []
+    print(f"Preparing {len(base_data)} items for judging...")
+
+    for k, base_item in base_data.items():
+        roi_item = roi_data[k]
+
+        # Optimization: Skip identical strings (Result is definitely same, so neither reg nor imp)
+        if str(base_item.get('response', '')).strip() == str(roi_item.get('response', '')).strip():
+            continue
+
+        judge_items.append({
+            'unique_key': k,
+            'dataset': base_item.get('dataset', 'unknown'),
+            'line_number': base_item.get('line_number', 'unknown'),
+            'id': base_item['id'],
+            'image': base_item['image'],
+            'question': base_item['prompt'],
+            'ground_truth': base_item['text'],
+            'base_resp': base_item.get('response', ''),
+            'roi_resp': roi_item.get('response', ''),
+            'surrounding_bbox': roi_item.get('surrounding_bbox'),
+            'roi_mask': roi_item.get('roi_mask')
+        })
+
+    print(f"Items requiring judgement (non-identical): {len(judge_items)}")
+
+    # 3. Parallel Processing
+    try:
+        mp.set_start_method('spawn', force=True)
+    except RuntimeError:
+        pass
+
+    gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
+    num_gpus = len(gpu_ids)
+
+    chunk_size = math.ceil(len(judge_items) / num_gpus)
+    chunks = [judge_items[i:i + chunk_size] for i in range(0, len(judge_items), chunk_size)]
+
+    manager = mp.Manager()
+    return_dict = manager.dict()
+    processes = []
+
+    print(f"Starting judging on {num_gpus} GPUs...")
+    for i in range(num_gpus):
+        chunk_data = chunks[i] if i < len(chunks) else []
+        p = mp.Process(target=worker_wrapper, args=(chunk_data, gpu_ids[i], args, return_dict))
+        p.start()
+        processes.append(p)
+
+    for p in processes:
+        p.join()
+
+    # 4. Aggregation
+    all_regressions = []
+    all_improvements = []
+
+    for gpu_id, res in return_dict.items():
+        all_regressions.extend(res['regressions'])
+        all_improvements.extend(res['improvements'])
+
+    print(f"\n✅ Analysis Complete.")
+    print(f"  - Regressions (Worse): {len(all_regressions)}")
+    print(f"  - Improvements (Better): {len(all_improvements)}")
+
+    # Save Regressions
+    with open(args.save_path, 'w', encoding='utf-8') as f:
+        json.dump(all_regressions, f, indent=2)
+    print(f"Saved regressions to {args.save_path}")
+
+    # Save Improvements
+    with open(args.save_path_improvements, 'w', encoding='utf-8') as f:
+        json.dump(all_improvements, f, indent=2)
+    print(f"Saved improvements to {args.save_path_improvements}")
+
+
+if __name__ == "__main__":
+    main()
+
+'''
+python analysis/filter_roi_fails.py \
+    --model-path /root/public/Qwen/Qwen2.5-VL-7B-Instruct \
+    --data-dir /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining  \
+    --save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2-short.json \
+    --save-path-improvements /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-short.json \
+    --batch-size 8 \
+    --gpu-ids 0,1
+
+python analysis/filter_roi_fails.py \
+    --model-path /root/public/Qwen/Qwen2.5-VL-7B-Instruct \
+    --data-dir /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining  \
+    --search-pattern "*_qwen2.5vl-3B.pkl" \
+    --save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-3b.json \
+    --save-path-improvements /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-3b.json \
+    --batch-size 8 \
+    --gpu-ids 0,1,2,3
+    
+'''

QWENVL-PRIVATE/analysis/fix_stage1_image_paths.py

@@ -0,0 +1,166 @@
+import argparse
+import os
+import pickle
+import re
+from pathlib import PureWindowsPath
+
+
+def normalize_path(path_str):
+    return path_str.replace("\\", "/")
+
+
+def is_abs_like(path_str):
+    if os.path.isabs(path_str):
+        return True
+    if re.match(r"^[A-Za-z]:\\\\", path_str):
+        return True
+    if re.match(r"^[A-Za-z]:/", path_str):
+        return True
+    return False
+
+
+def maybe_make_relative(path_str, prefixes, fallback_basename=False):
+    if not isinstance(path_str, str):
+        return path_str, False
+
+    if not is_abs_like(path_str):
+        return path_str, False
+
+    raw = path_str
+    norm = normalize_path(path_str).rstrip("/")
+    for prefix in prefixes:
+        if not prefix:
+            continue
+        prefix_norm = normalize_path(prefix).rstrip("/")
+        if norm == prefix_norm:
+            return os.path.basename(norm), True
+        if norm.startswith(prefix_norm + "/"):
+            return norm[len(prefix_norm) + 1 :], True
+
+    if fallback_basename:
+        return os.path.basename(norm), True
+
+    return raw, False
+
+
+def iter_item_dicts(obj):
+    if isinstance(obj, list):
+        for item in obj:
+            if isinstance(item, dict):
+                yield item
+    elif isinstance(obj, dict):
+        # Common cases: dict of items or dict with list values
+        for value in obj.values():
+            if isinstance(value, dict):
+                yield value
+            elif isinstance(value, list):
+                for item in value:
+                    if isinstance(item, dict):
+                        yield item
+
+
+def process_pkl(path, prefixes, fallback_basename=False):
+    with open(path, "rb") as f:
+        try:
+            data = pickle.load(f)
+        except EOFError:
+            print(f"Warning: empty or corrupted pickle, skipping: {path}")
+            return None, 0, 0
+
+    updated = 0
+    total = 0
+    for item in iter_item_dicts(data):
+        if "image" not in item:
+            continue
+        total += 1
+        new_val, changed = maybe_make_relative(
+            item["image"], prefixes, fallback_basename=fallback_basename
+        )
+        if changed:
+            item["image"] = new_val
+            updated += 1
+
+    return data, total, updated
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Replace absolute image paths with relative paths in stage1 pkls."
+    )
+    parser.add_argument(
+        "--input-dir",
+        type=str,
+        required=True,
+        help="Directory containing .pkl files to process.",
+    )
+    parser.add_argument(
+        "--prefix-to-strip",
+        type=str,
+        action="append",
+        default=[],
+        help="Prefix to strip from absolute image paths. Can be set multiple times.",
+    )
+    parser.add_argument(
+        "--dataset-root",
+        type=str,
+        default="",
+        help="Optional dataset root path to strip (added to prefix list if set).",
+    )
+    parser.add_argument(
+        "--fallback-basename",
+        action="store_true",
+        help="If no prefix matches, fallback to basename for absolute paths.",
+    )
+    parser.add_argument(
+        "--out-dir",
+        type=str,
+        default="",
+        help="Output directory for fixed pkls (default: input dir).",
+    )
+    parser.add_argument(
+        "--inplace",
+        action="store_true",
+        help="Overwrite pkls in place (ignores --out-dir).",
+    )
+    args = parser.parse_args()
+
+    prefixes = list(args.prefix_to_strip)
+    if args.dataset_root:
+        prefixes.append(args.dataset_root)
+
+    input_dir = os.path.abspath(args.input_dir)
+    out_dir = input_dir if args.inplace else (args.out_dir or input_dir)
+    os.makedirs(out_dir, exist_ok=True)
+
+    pkl_files = [
+        os.path.join(input_dir, name)
+        for name in os.listdir(input_dir)
+        if name.endswith(".pkl")
+    ]
+
+    if not pkl_files:
+        print(f"No .pkl files found in {input_dir}")
+        return
+
+    for pkl_path in pkl_files:
+        data, total, updated = process_pkl(
+            pkl_path, prefixes, fallback_basename=args.fallback_basename
+        )
+        if data is None:
+            continue
+        if args.inplace:
+            out_path = pkl_path
+        else:
+            base = os.path.basename(pkl_path)
+            out_path = os.path.join(out_dir, base.replace(".pkl", "_relpath.pkl"))
+
+        with open(out_path, "wb") as f:
+            pickle.dump(data, f)
+
+        print(
+            f"{os.path.basename(pkl_path)}: updated {updated}/{total} image entries -> {out_path}"
+        )
+
+
+if __name__ == "__main__":
+    main()

QWENVL-PRIVATE/analysis/gated_roi_analysis.py

@@ -0,0 +1,899 @@
+import sys
+import os
+script_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(script_dir)
+sys.path.append(project_root)
+
+from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
+from qwen_src.qwen2_vl import Qwen2VLForConditionalGeneration
+import torch
+import matplotlib.pyplot as plt
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor
+import math
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map, plot_attention_across_layers, \
+ plot_tensor_2d, visualize_token_wise_attention,get_visual_feat_norm, visualize_norm_values, plot_attention_analysis
+import json
+import os
+from tqdm import tqdm
+import torchvision.transforms.functional as torchvision_F
+from qwen_src.mm_utils import expand2square, create_pseudo_labels
+import matplotlib.patches as patches
+from matplotlib.colors import ListedColormap, BoundaryNorm
+import pickle
+import argparse
+import multiprocessing as mp
+import pandas as pd  # <-- Make sure to import pandas
+from matching_utils import TextVQAAccuracyEvaluator
+from qwen_vl_utils import process_vision_info
+import requests
+import io
+
+def get_KL_divergence(logits_list,):
+    import torch.nn as nn
+    from torch.nn import functional as F
+    KL_div = nn.KLDivLoss(reduction='batchmean')
+    KL_div_list = [KL_div(F.log_softmax(logits_list[-1], dim=1), F.softmax(logits_list[i], dim=1)).item() for i in range(len(logits_list))]
+    return KL_div_list
+
+
+def plot_kl_divergence_distribution(kl_div_pickle_file,
+                                    save_path,
+                                    trim_threshold=None,
+                                    bins=50,
+                                    title="Distribution of KL Divergence"):
+    print(f"Loading KL data from {kl_div_pickle_file}...")
+    try:
+        with open(kl_div_pickle_file, 'rb') as f:
+            KL_div_dict = pickle.load(f)
+    except FileNotFoundError:
+        print(f"Error: Pickle file not found at {kl_div_pickle_file}")
+        return
+    except Exception as e:
+        print(f"Error loading pickle file: {e}")
+        return
+
+    # --- 1. Load and process all data ---
+    processed_data = []
+    for line_num, s in KL_div_dict.items():
+        if 'KL_divergence' in s and 'base_answer' in s and 'two_stage_answer' in s:
+            processed_data.append({
+                'kl': s['KL_divergence'],
+                'different': str(s['base_answer']) != str(s['two_stage_answer'])
+            })
+        else:
+            print(f"Warning: Skipping item {line_num}, missing required keys.")
+
+    if not processed_data:
+        print("No valid data found in the pickle file.")
+        return
+
+    df = pd.DataFrame(processed_data)
+    original_count = len(df)
+    plot_title = title
+
+    # --- 2. Handle trimming (clipping) ---
+    # Apply your logic: clip values larger than the threshold
+    if trim_threshold is not None:
+        kl_for_hist = df['kl'].clip(upper=trim_threshold)
+        print(f"Original data points: {original_count}")
+        plot_title = f"{title}\n(Values > {trim_threshold} clipped to {trim_threshold})"
+    else:
+        kl_for_hist = df['kl']
+        print(f"Plotting all {original_count} data points. No trimming.")
+
+    # --- 3. Plot 1: KL Divergence Histogram ---
+    print("Generating Histogram plot...")
+    plt.figure(figsize=(10, 6))
+
+    # Create the histogram AND capture the bin edges
+    counts, bin_edges, _ = plt.hist(kl_for_hist, bins=bins,
+                                    edgecolor='k', alpha=0.75, color='skyblue')
+
+    # Add a vertical line for the mean
+    mean_val = np.mean(kl_for_hist)
+    plt.axvline(mean_val, color='r', linestyle='dashed', linewidth=2)
+
+    # Get current y-axis limits to position text
+    ymin, ymax = plt.ylim()
+    plt.text(mean_val * 1.05, ymax * 0.9, f'Mean: {mean_val:.3f}', color='r')
+
+    plt.xlabel("KL Divergence Score")
+    plt.ylabel("Frequency (Count)")
+    plt.title(plot_title)
+    plt.grid(True, linestyle='--', alpha=0.6)
+    plt.tight_layout()
+    plt.show()
+    plt.close()  # Close the plot to free memory
+
+    # --- 4. Calculate Ratios for Second Plot ---
+    print("Calculating ratios for the second plot...")
+
+    # Use pandas.cut to put each sample into the bins created by plt.hist
+    # We use the *clipped* KL values to match the histogram
+    df['bin'] = pd.cut(kl_for_hist, bins=bin_edges,
+                       include_lowest=True, right=True)
+
+    # Group by the bins and calculate the mean of 'different' (True=1, False=0)
+    bin_ratios = df.groupby('bin')['different'].mean()
+
+    # Get the center of each bin for plotting
+    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
+
+    # --- 5. Plot 2: Answer Difference Ratio ---
+    print("Generating Answer Difference Ratio plot...")
+    plt.figure(figsize=(10, 6))
+
+    plt.bar(bin_centers, bin_ratios.fillna(0),
+            width=np.diff(bin_edges),
+            edgecolor='k', alpha=0.85, color='salmon')
+
+    plt.xlabel("KL Divergence Bin")
+    plt.ylabel("Ratio of Different Answers")
+    plt.title("Answer Difference Ratio per KL Divergence Bin")
+    plt.grid(True, linestyle='--', alpha=0.6)
+    plt.ylim(0, 1.0)  # Ratio is always between 0 and 1
+    plt.tight_layout()
+    plt.show()
+    plt.close()
+def case_study_KL_divergence(model, processor,
+                             question_file, image_folder,
+                             kl_div_pickle_file, threshold,
+                             max_cases=10, max_new_tokens=1024):
+    """
+    Performs a case study on questions with high KL divergence between
+    base and two-stage model logits.
+
+    Args:
+        model: The multimodal model.
+        processor: The model's processor.
+        question_file (str): Path to the JSONL file with all questions.
+        image_folder (str): Path to the folder containing images.
+        kl_div_pickle_file (str): Path to the .pkl file containing the
+                                  dictionary of {question_id: kl_divergence}.
+        threshold (float): The minimum KL divergence score to include
+                           in the case study.
+        max_cases (int, optional): The maximum number of cases to print.
+                                   Defaults to 10.
+        max_new_tokens (int, optional): The max tokens for generation.
+                                        Defaults to 1024.
+    """
+
+    print(f"🚀 Starting case study for KL Divergence > {threshold}")
+
+    # 1. Load KL divergence data
+    print(f"Loading KL divergence data from: {kl_div_pickle_file}")
+    with open(kl_div_pickle_file, 'rb') as f:
+        KL_div_dict = pickle.load(f)
+
+    # 2. Load and preprocess all questions
+    questions = [json.loads(q) for q in open(os.path.expanduser(question_file), "r")]
+
+    questions_dict = {}
+    for tmp_ques in questions:
+        qs = tmp_ques['text']
+        # Remove lines with 'Reference OCR' as in the original function
+        if 'Reference OCR' in qs:
+            qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
+        tmp_ques['text'] = qs
+        questions_dict[tmp_ques['question_id']] = tmp_ques
+
+    # 3. Filter for high KL divergence questions
+    high_kl_cases = []
+    for q_id, kl_div_sample in KL_div_dict.items():
+        kl_div = kl_div_sample['KL_divergence']
+        if kl_div > threshold:
+            if q_id in questions_dict:
+                high_kl_cases.append((q_id, kl_div, questions_dict[q_id]))
+            else:
+                print(f"Warning: Question ID {q_id} from pickle file not found in question file.")
+
+    # # Sort by KL divergence (highest first)
+    # high_kl_cases.sort(key=lambda x: x[1], reverse=True)
+
+    if not high_kl_cases:
+        print("No questions found above the threshold.")
+        return
+
+    # Set processor image size limits (matching your eval function)
+    processor.image_processor.max_pixels = 576 * 28 * 28
+    processor.image_processor.min_pixels = 576 * 28 * 28
+
+    # 4. Loop through filtered cases and generate answers
+    print(f"Found {len(high_kl_cases)} cases above threshold. Analyzing top {min(len(high_kl_cases), max_cases)}...")
+
+    model.model.is_debug = True
+    case_count = 0
+    for q_id, kl_value, line in high_kl_cases:
+        if case_count >= max_cases:
+            print(f"Reached max_cases limit of {max_cases}.")
+            break
+
+        case_count += 1
+
+        try:
+            image_file = line["image"]
+            qs = "<image>\n" + line["text"]
+            image_path = os.path.join(image_folder, image_file)
+            image = Image.open(image_path).convert('RGB')
+
+            # Prepare inputs
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": f"{qs}"},
+                        {"image": image_path}  # Note: path is used by template, image obj by processor
+                    ]
+                }
+            ]
+            text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '').replace('<image>',
+                                                                                           '<image><|vision_start|><|image_pad|><|vision_end|>')
+            inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt").to('cuda')
+            inputs.data['src_images'] = [image]
+            inputs.data['processor'] = processor
+            input_ids_len = inputs.input_ids.shape[1]
+
+            print(f"\n--- CASE STUDY {case_count} (Q_ID: {q_id}) ---")
+            print(f"Image File: {image_file}")
+            print(f"KL Divergence: {kl_value:.4f}")
+            print(f"Question: {line['text']}")
+            print("---")
+
+            # --- 4a. Generate with two_stage_flag = False (Base Model) ---
+            two_stage_flag = False
+            model.model.enable_twig = two_stage_flag
+            model.model.roi_enable2stage = two_stage_flag
+            model.model.two_stage_vanilla = two_stage_flag
+            model.roi_enable2stage = two_stage_flag
+            model.model.roi_conf_thresh = 0.05
+
+            with torch.no_grad():
+                output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                            pad_token_id=processor.tokenizer.eos_token_id)
+
+            generated_ids = output_ids[:, input_ids_len:]
+            base_answer = \
+            processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
+            print(f"**Base Answer (RoI OFF):**\n{base_answer}\n")
+
+            # --- 4b. Generate with two_stage_flag = True (Two-Stage Model) ---
+            two_stage_flag = True
+            model.model.enable_twig = two_stage_flag
+            model.model.roi_enable2stage = two_stage_flag
+            model.model.two_stage_vanilla = two_stage_flag
+            model.roi_enable2stage = two_stage_flag
+            model.model.roi_conf_thresh = 0.05
+
+            with torch.no_grad():
+                output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                            pad_token_id=processor.tokenizer.eos_token_id)
+
+            generated_ids = output_ids[:, input_ids_len:]
+            two_stage_answer = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
+            print(f"**Two-Stage Answer (RoI ON):**\n{two_stage_answer}\n")
+            print("-" * 40)
+
+        except Exception as e:
+            print(f"Error processing question {q_id}: {e}")
+            continue
+
+    print("✅ Case study complete.")
+
+
+def eval_roi_KLdivergence(model, processor, question_file, image_folder, dst_save_path, max_new_tokens=1024):
+
+    questions = [json.loads(q) for q in open(os.path.expanduser(question_file), "r")]
+    system_prompt= "You are a helpful assistant"
+    sample = 10
+    two_stage_flag = False
+
+    processor.image_processor.max_pixels = 576 * 28 * 28
+    processor.image_processor.min_pixels = 576 * 28 * 28
+
+    new_questions = []
+    for tmp_ques in questions:
+        qs = tmp_ques['text']
+        if 'Reference OCR' in qs:
+            # Remove lines with 'Reference OCR
+            qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
+        tmp_ques['text'] = qs
+        new_questions.append(tmp_ques)
+    questions = new_questions
+    model.model.is_debug = False
+    KL_div_analysis_dict = {}
+    for line in tqdm(questions):
+        image_file = line["image"]
+        qs = "<image>\n" + line["text"]
+        image = Image.open(os.path.join(image_folder, image_file)).convert('RGB')
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": f"{qs}"
+                    },
+                    {
+                        "image": os.path.join(image_folder, image_file)
+                    }
+                ]
+            }
+        ]
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '').replace('<image>', '<image><|vision_start|><|image_pad|><|vision_end|>')
+        inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt").to('cuda')
+        inputs.data['src_images'] = [image]
+        inputs.data['processor'] = processor
+        input_ids_len = inputs.input_ids.shape[1]
+
+        two_stage_flag = False
+        model.model.enable_twig = two_stage_flag
+        model.model.roi_enable2stage = two_stage_flag
+        model.model.two_stage_vanilla = two_stage_flag
+        model.roi_enable2stage = two_stage_flag
+        model.model.roi_conf_thresh = 0.05
+        with torch.no_grad():
+            output_dicts = model.forward(**inputs.data,
+                                         output_attentions=True,
+                                         return_dict=False,
+                                         output_hidden_states=False
+                                         )
+            base_logits = output_dicts[0][:,-1]
+        with torch.no_grad():
+            output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                        pad_token_id=processor.tokenizer.eos_token_id)
+
+        generated_ids = output_ids[:, input_ids_len:]
+        base_answer = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
+
+        two_stage_flag = True
+        model.model.enable_twig = two_stage_flag
+        model.model.roi_enable2stage = two_stage_flag
+        model.model.two_stage_vanilla = two_stage_flag
+        model.roi_enable2stage = two_stage_flag
+        model.model.roi_conf_thresh = 0.05
+        with torch.no_grad():
+            output_dicts = model.forward(**inputs.data,
+                                         output_attentions=False,
+                                         return_dict=False,
+                                         output_hidden_states=False
+                                         )
+            two_stage_logits = output_dicts[0][:, -1]
+
+        logits_list = [base_logits, two_stage_logits]
+        KL_div_list = get_KL_divergence(logits_list)
+
+        with torch.no_grad():
+            output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                        pad_token_id=processor.tokenizer.eos_token_id)
+
+        generated_ids = output_ids[:, input_ids_len:]
+        two_stage_answer = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
+        print(f"**Two-Stage Answer (RoI ON):**\n{two_stage_answer}\n")
+
+        print(f"\n--- Processing sample: {image_file} ---")
+        print(f"Question: {qs}")
+        print("KL Divergence between RoI on and off:", KL_div_list[0])
+        sample += 1
+        tmp_sample = {
+            'image_file': image_file,
+            'question': qs,
+            'base_answer': base_answer,
+            'two_stage_answer': two_stage_answer,
+            'KL_divergence': KL_div_list[0]
+        }
+        KL_div_analysis_dict[line['question_id']] = tmp_sample
+        pass
+
+    print(f'saving KL divergence analysis dict to:{dst_save_path}')
+    with open(dst_save_path, 'wb') as f:
+        pickle.dump(KL_div_analysis_dict, f)
+
+
+def worker_process_kl_divergence(questions_chunk, gpu_id, args):
+    """
+    This is the worker function that runs on a single GPU.
+    It loads the model and processes its assigned chunk of questions.
+    """
+
+    device = f'cuda:{gpu_id}'
+    KL_div_analysis_dict = {}
+    max_new_tokens = 1024
+    if not questions_chunk:
+        print(f"Worker {gpu_id}: No questions in chunk. Exiting.")
+        return KL_div_analysis_dict
+
+    try:
+        # 1. Load Model and Processor INSIDE the worker
+        print(f"Worker {gpu_id}: Loading processor from {args.model_path}...")
+        processor = AutoProcessor.from_pretrained(
+            'Qwen/Qwen2.5-VL-7B-Instruct',
+            trust_remote_code=True,
+            padding_side='left', use_fast=True
+        )
+
+        print(f"Worker {gpu_id}: Loading model from {args.model_path}...")
+        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device  # Assign model directly to the correct device
+        )
+        model.eval()
+
+        # Set processor image size limits (matching your original function)
+        processor.image_processor.max_pixels = 576 * 28 * 28
+        processor.image_processor.min_pixels = 576 * 28 * 28
+
+        # 2. Process the chunk of questions
+        print(f"Worker {gpu_id}: Starting processing of {len(questions_chunk)} items.")
+
+        # Use tqdm for progress, disable for all but the first worker if too noisy
+        progress_bar = tqdm(questions_chunk, desc=f"GPU {gpu_id}", position=gpu_id)
+
+        for line in progress_bar:
+            try:
+                # --- MODIFIED: Handle different dataset formats ---
+                image_file_str = ''
+                ground_truth = None
+
+                if args.question_file is not None:  # Case 1: Old format with "text" key
+                    qs_text = line["text"]
+                    qs_text = qs_text + " Answer the question using a single word or phrase." if "Answer the question using a single word or phrase." not in qs_text else qs_text
+                    image_file = line["image"]
+                    image_path = os.path.join(args.image_folder, image_file)
+                    image = Image.open(image_path).convert('RGB')
+                    image_file_str = image_file
+
+                elif args.hf_dataset is not None:  # Case 2: New HF Dataset format (e.g., vstar_bench)
+                    if 'vstar' in args.hf_dataset:
+                        qs_text = line["text"]
+                        qs_text = qs_text + " Answer the question using a single word or phrase." if "Answer the question using a single word or phrase." not in qs_text else qs_text
+                        image = line["image"].convert('RGB')
+                        ground_truth = line.get("label")  # Store the ground truth
+                        image_file_str = f"hf_dataset_{line.get('question_id', 'unknown')}"
+                    elif 'POPE' in args.hf_dataset:
+                        qs_text = line["question"] + " Answer the question using a single word or phrase."
+                        image = line["image"].convert('RGB')
+                        ground_truth = line.get("answer")
+                        image_file_str = f"hf_dataset_{line.get('question_id', 'unknown')}"
+
+                else:
+                    print(f"Worker {gpu_id}: Skipping line {line.get('original_line_number')}, unknown format.")
+                    continue
+                # --- END MODIFICATION ---
+                qs = "<image>\n" + qs_text
+                messages = [
+                    {"role": "user", "content": [{"type": "text", "text": f"{qs}"}, {"image": image}]}
+                ]
+                text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+                text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '').replace('<image>',
+                                                                                               '<image><|vision_start|><|image_pad|><|vision_end|>')
+
+                inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt").to(device)
+
+                # We don't need these for KL div, but keeping if your model.forward needs them
+                inputs.data['src_images'] = [image]
+                inputs.data['processor'] = processor
+                input_ids_len = inputs.input_ids.shape[1]
+
+                # --- Run with two_stage_flag = False (Base) ---
+                two_stage_flag = False
+                model.model.enable_twig = two_stage_flag
+                model.model.roi_enable2stage = two_stage_flag
+                model.model.two_stage_vanilla = two_stage_flag
+                model.roi_enable2stage = two_stage_flag
+                model.model.roi_conf_thresh = 0.05
+                with torch.no_grad():
+                    output_dicts = model.forward(**inputs.data,
+                                                 output_attentions=True,
+                                                 return_dict=False,
+                                                 output_hidden_states=False
+                                                 )
+                    base_logits = output_dicts[0][:, -1]
+                with torch.no_grad():
+                    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                                pad_token_id=processor.tokenizer.eos_token_id)
+
+                generated_ids = output_ids[:, input_ids_len:]
+                base_answer = \
+                    processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[
+                        0]
+
+                # --- Run with two_stage_flag = True (Two-Stage) ---
+                two_stage_flag = True
+                model.model.enable_twig = two_stage_flag
+                model.model.roi_enable2stage = two_stage_flag
+                model.model.two_stage_vanilla = two_stage_flag
+                model.roi_enable2stage = two_stage_flag
+                model.model.roi_conf_thresh = 0.05
+                with torch.no_grad():
+                    output_dicts = model.forward(**inputs.data,
+                                                 output_attentions=False,
+                                                 return_dict=False,
+                                                 output_hidden_states=False
+                                                 )
+                    two_stage_logits = output_dicts[0][:, -1]
+
+                logits_list = [base_logits, two_stage_logits]
+                KL_div_list = get_KL_divergence(logits_list)
+
+                with torch.no_grad():
+                    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                                pad_token_id=processor.tokenizer.eos_token_id)
+
+                generated_ids = output_ids[:, input_ids_len:]
+                two_stage_answer = \
+                    processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[
+                        0]
+
+                tmp_sample = {
+                    'image_file': image_file_str,
+                    'question': qs_text,  # Store the clean question
+                    'base_answer': base_answer,
+                    'two_stage_answer': two_stage_answer,
+                    'KL_divergence': KL_div_list[0],
+                    'ground_truth': ground_truth  # This will be None for old data, or the answer
+                }
+
+                # *** MODIFIED LINE ***
+                # Use the original_line_number as the unique key
+                line_number = line['original_line_number']
+                if line_number in KL_div_analysis_dict:
+                    print(
+                        f"Worker {gpu_id}: Warning - Duplicate line_number {line_number} found. Overwriting previous entry.")
+                KL_div_analysis_dict[line_number] = tmp_sample
+
+            except Exception as e:
+                # *** MODIFIED LINE ***
+                # Update error logging to use line_number
+                print(f"Worker {gpu_id}: Error processing {line.get('original_line_number', 'unknown line')}: {e}")
+                continue
+
+        progress_bar.close()
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: A critical error occurred: {e}")
+    finally:
+        # Clean up model from GPU memory
+        del model
+        del processor
+        torch.cuda.empty_cache()
+        print(f"Worker {gpu_id}: Finished and cleaned up.")
+
+    return KL_div_analysis_dict
+
+
+def evaluate_accuracy_with_kl_bins(kl_div_pickle_file,
+                                   plot_save_path,
+                                   annotation_file=None,  # <-- Made optional
+                                   trim_threshold=None,
+                                   bins=20):
+    """
+    Evaluates accuracy for base and two-stage answers, and plots
+    accuracy vs. KL divergence bins.
+
+    If annotation_file is provided, it's used as the GT source.
+    If None, 'ground_truth' key from the pickle file is used.
+    """
+
+    gt_answers_by_line = {}
+    use_external_gt = False
+
+    # --- 1. Load Ground Truth (Conditionally) ---
+    if annotation_file:
+        use_external_gt = True
+        print(f"Loading external annotations from {annotation_file}...")
+        try:
+            # Your logic for loading the .json file's 'data' list
+            if annotation_file.endswith('.jsonl'):
+                annotations = [json.loads(line) for line in open(os.path.expanduser(annotation_file), 'r')]
+            else:
+                annotations = json.load(open(os.path.expanduser(annotation_file)))['data']
+            gt_answers_by_line = {i: ann['answers'] for i, ann in enumerate(annotations)}
+            if not gt_answers_by_line:
+                print(f"Warning: No annotations loaded from {annotation_file}.")
+        except Exception as e:
+            print(f"Error loading annotation file: {e}")
+            return
+    else:
+        print("No annotation file provided. Using 'ground_truth' from pickle file.")
+
+    # --- 2. Load Prediction Data (from your pickle file) ---
+    print(f"Loading KL data from {kl_div_pickle_file}...")
+    try:
+        with open(kl_div_pickle_file, 'rb') as f:
+            KL_div_dict = pickle.load(f)
+    except FileNotFoundError:
+        print(f"Error: Pickle file not found at {kl_div_pickle_file}")
+        return
+    except Exception as e:
+        print(f"Error loading pickle file: {e}")
+        return
+
+    # --- 3. Match Data and Calculate Single-Sample Accuracy ---
+    print("Matching predictions to ground truth and calculating scores...")
+    evaluator = TextVQAAccuracyEvaluator()  # Use the same evaluator for both paths
+    data_for_df = []
+
+    for line_number, sample in tqdm(KL_div_dict.items(), desc="Evaluating samples"):
+
+        gt_answers = None
+
+        if use_external_gt:
+            # Path 1: Get GT from the external file
+            if line_number not in gt_answers_by_line:
+                print(f"Warning: Line {line_number} from pickle not found in annotations. Skipping.")
+                continue
+            gt_answers = gt_answers_by_line[line_number]
+        else:
+            # Path 2: Get GT from the pickle file itself
+            gt = sample.get('ground_truth') if 'ground_truth' in sample else sample['text']
+            if gt is None:
+                # print(f"Warning: Line {line_number} in pickle has no 'ground_truth' key. Skipping.")
+                continue
+
+            # Ensure gt_answers is a list, as the evaluator expects
+            gt_answers = [gt.lower()] if isinstance(gt, str) else gt
+
+        if gt_answers is None:
+            continue
+
+        # We can now use the same evaluation logic for both paths
+        if annotation_file is not None:
+            base_acc = evaluator.eval_pred_list([
+                {"pred_answer": sample['base_answer'], "gt_answers": gt_answers}
+            ])
+
+            two_stage_acc = evaluator.eval_pred_list([
+                {"pred_answer": sample['two_stage_answer'], "gt_answers": gt_answers}
+            ])
+        else:
+            base_acc = sample['base_answer'].lower() in gt_answers
+            two_stage_acc = sample['two_stage_answer'].lower() in gt_answers
+        data_for_df.append({
+            'kl': sample['KL_divergence'],
+            'base_acc': base_acc,
+            'two_stage_acc': two_stage_acc
+        })
+
+    if not data_for_df:
+        print("No matching data found to evaluate. Exiting.")
+        return
+
+    df = pd.DataFrame(data_for_df)
+
+    # --- 4. Print Overall Accuracy (Unchanged) ---
+    overall_base_acc = df['base_acc'].mean() * 100.0
+    overall_two_stage_acc = df['two_stage_acc'].mean() * 100.0
+
+    print("\n" + "=" * 30)
+    print("📊 Overall Accuracy Results:")
+    print(f"  Base Answer Accuracy:     {overall_base_acc:.2f}%")
+    print(f"  Two-Stage Answer Accuracy: {overall_two_stage_acc:.2f}%")
+    print("=" * 30 + "\n")
+
+    # --- 5. Bin Data and Prepare for Plotting (Unchanged) ---
+    print("Binning data and preparing plot...")
+
+    if trim_threshold is not None:
+        kl_for_binning = df['kl'].clip(upper=trim_threshold)
+        plot_title = f"Accuracy vs. KL Divergence (Values > {trim_threshold} clipped)"
+    else:
+        kl_for_binning = df['kl']
+        plot_title = "Accuracy vs. KL Divergence"
+
+    df['bin'] = pd.cut(kl_for_binning, bins=bins, include_lowest=True, right=True)
+
+    binned_acc = df.groupby('bin', observed=False)[['base_acc', 'two_stage_acc']].mean()
+    binned_acc = binned_acc.fillna(0)
+
+    bin_centers = (binned_acc.index.categories.left + binned_acc.index.categories.right) / 2
+    bin_widths = (binned_acc.index.categories.right - binned_acc.index.categories.left)
+    bar_width = bin_widths * 0.4
+
+    # --- 6. Plot the Grouped Bar Chart (Unchanged, except save path) ---
+    plt.figure(figsize=(15, 7))
+
+    plt.bar(bin_centers - bar_width / 2,
+            binned_acc['base_acc'],
+            width=bar_width,
+            label='Base Accuracy',
+            color='skyblue',
+            edgecolor='k')
+
+    plt.bar(bin_centers + bar_width / 2,
+            binned_acc['two_stage_acc'],
+            width=bar_width,
+            label='Two-Stage Accuracy',
+            color='salmon',
+            edgecolor='k')
+
+    plt.xlabel("KL Divergence Bin")
+    plt.ylabel("Mean Accuracy")
+    plt.title(plot_title)
+    plt.legend()
+    plt.grid(True, linestyle='--', alpha=0.6, axis='y')
+    plt.xticks(bin_centers, [f"{b.left:.2f}" for b in binned_acc.index.categories], rotation=45, ha='right')
+    plt.ylim(0, 1.0)
+    plt.tight_layout()
+
+    # try:
+    #     plt.savefig(plot_save_path)
+    #     print(f"✅ Accuracy plot saved successfully to {plot_save_path}")
+    # except Exception as e:
+    #     print(f"Error saving plot: {e}")
+
+    plt.show()
+    plt.close()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Parallel KL Divergence Evaluation")
+    parser.add_argument('--model-path', type=str, default=None,
+                        help="Path to the model and processor")
+    parser.add_argument('--question-file', type=str, default=None,
+                        help="Path to the JSONL question file (E.g., TextVQA)")
+    parser.add_argument('--image-folder', type=str, default=None,
+                        help="Path to the image directory (for JSONL file)")
+    parser.add_argument('--dst-save-path', type=str, required=True,
+                        help="Path to save the output .pkl file")
+    parser.add_argument('--gpu-ids', type=str, default='0,1',
+                        help="Comma-separated list of GPU IDs to use (e.g., '0,1,2,3')")
+    # --- ADDED: New HF Dataset args ---
+    parser.add_argument('--hf_dataset', type=str, default=None,
+                        help="Name of the Hugging Face dataset (e.g., 'craigwu/vstar_bench')")
+    parser.add_argument('--dataset_split', type=str, default='test',
+                        help="Split to use from the Hugging Face dataset (e.g., 'test')")
+    ### --- END ADDITION ---
+    args = parser.parse_args()
+
+    # # 1. Prepare Data and Split into Chunks
+    # gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
+    # num_gpus = len(gpu_ids)
+    # questions = []
+    #
+    # if args.hf_dataset:
+    #     print(f"Loading Hugging Face dataset {args.hf_dataset}, split {args.dataset_split}")
+    #     # Import only if needed
+    #     from datasets import load_dataset
+    #     from datasets import Image as HFImage
+    #
+    #     dataset = load_dataset(args.hf_dataset, split=args.dataset_split)
+    #     print("Casting 'image' column to load from cache...")
+    #     questions_list = list(dataset)[:3000]
+    #
+    #     new_questions = []
+    #     for i, tmp_ques in enumerate(questions_list):
+    #         tmp_ques['original_line_number'] = i
+    #         new_questions.append(tmp_ques)
+    #     questions = new_questions
+    #     print(f"Loaded {len(questions)} preprocessed questions from Hugging Face.")
+    #
+    # elif args.question_file:
+    #     print("Loading and preprocessing questions from JSONL file...")
+    #     if not args.image_folder:
+    #         raise ValueError("--image-folder is required when using --question-file")
+    #
+    #     questions_list = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")][:3000]
+    #
+    #     new_questions = []
+    #     for i, tmp_ques in enumerate(questions_list):
+    #         qs = tmp_ques['text']
+    #         if 'Reference OCR' in qs:
+    #             qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
+    #         tmp_ques['text'] = qs
+    #         tmp_ques['original_line_number'] = i  # Add the line number as a new field
+    #         new_questions.append(tmp_ques)
+    #     questions = new_questions
+    #     print(f"Loaded {len(questions)} preprocessed questions from local file.")
+    #
+    # else:
+    #     raise ValueError("You must provide either --question-file or --hf_dataset")
+    # chunk_size = math.ceil(len(questions) / num_gpus)
+    # chunks = [questions[i:i + chunk_size] for i in range(0, len(questions), chunk_size)]
+    # # 2. Set up arguments for each worker process
+    # worker_args = []
+    # for i in range(num_gpus):
+    #     chunk_data = chunks[i] if i < len(chunks) else []
+    #     worker_args.append((chunk_data, gpu_ids[i], args))
+    #
+    # # 3. Run processes in parallel
+    # print(f"Starting parallel processing on {num_gpus} GPUs: {gpu_ids}")
+    # # Use 'spawn' start method for CUDA compatibility
+    # ctx = mp.get_context('spawn')
+    # with ctx.Pool(processes=num_gpus) as pool:
+    #     results_from_pool = pool.starmap(worker_process_kl_divergence, worker_args)
+    #
+    # # 4. Gather and Save Results
+    # print("All chunks processed. Gathering results...")
+    # final_KL_div_dict = {}
+    # for sub_dict in results_from_pool:
+    #     final_KL_div_dict.update(sub_dict)
+    #
+    # with open(os.path.expanduser(args.dst_save_path), 'wb') as f:
+    #     pickle.dump(final_KL_div_dict, f)
+    # print(f"✅ Successfully processed and saved {len(final_KL_div_dict)} KL divergence scores to {args.dst_save_path}.")
+
+    # #output_text, input_height, input_width = inference()
+    # use_act = False
+    # device = 'cuda'
+    # # # model_path = 'Qwen/Qwen2.5-VL-3B-Instruct'
+    # # # model_path = 'Qwen/Qwen2.5-VL-7B-Instruct'
+    # model_path = '/home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled'
+    # model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    #     model_path,
+    #     torch_dtype=torch.bfloat16,
+    #     attn_implementation="eager",
+    # ).eval().to(device)
+    # processor = AutoProcessor.from_pretrained('Qwen/Qwen2.5-VL-7B-Instruct', trust_remote_code=True, padding_side='left', use_fast=True)
+
+    # # # Run the case study on questions with KL divergence > 1.0
+    # case_study_KL_divergence(
+    #     model,
+    #     processor,
+    #     QUESTION_FILE,
+    #     IMAGE_FOLDER,
+    #     KL_PICKLE_PATH,
+    #     threshold=1.0,  # <-- Set your desired threshold here
+    #     max_cases=100  # <-- Set how many examples you want to see
+    # )
+
+
+    #KL_PICKLE_PATH = '../excluded/roi_KLdiv_textvqa_576res.pkl'
+    # QUESTION_FILE = '.../../llava/playground/data/eval/textvqa/llava_textvqa_val_v051_ocr.jsonl'
+    # IMAGE_FOLDER = '../../llava/playground/data/eval/textvqa/train_images'
+    ANNO_FILE = "/home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl" #'../../llava/playground/data/eval/textvqa/TextVQA_0.5.1_val.json'
+
+    KL_PICKLE_PATH = args.dst_save_path
+    ANNO_FILE = None if args.hf_dataset else ANNO_FILE
+
+    plot_kl_divergence_distribution(
+        kl_div_pickle_file=KL_PICKLE_PATH,
+        save_path=None,
+        trim_threshold=5.0,
+        bins=50,
+        title="Distribution of KL Divergence (RoI On vs Off) on Puseudo Train Set"
+    )
+
+    ###check accuracy vs KL divergence bins
+    evaluate_accuracy_with_kl_bins(
+        kl_div_pickle_file=KL_PICKLE_PATH,
+        annotation_file=None,
+        plot_save_path=None,
+        trim_threshold=5.0,
+        bins=50
+    )
+    pass
+
+'''
+CUDA_VISIBLE_DEVICES=1,3 python analysis/gated_roi_analysis.py \
+    --model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+    --question-file ../llava/playground/data/eval/textvqa/llava_textvqa_val_v051_ocr.jsonl \
+    --image-folder ../llava/playground/data/eval/textvqa/train_images \
+    --dst-save-path ./excluded/roi_KLdiv_textvqa_576res.pkl \
+    --gpu-ids 0,1
+    
+    
+CUDA_VISIBLE_DEVICES=0,1,2,3 python analysis/gated_roi_analysis.py \
+--model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+--hf_dataset "lmms-lab/vstar-bench" --dataset_split "test" \
+--dst-save-path ./excluded/vstar_analysis.pkl --gpu-ids 0,1,2,3
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 python analysis/gated_roi_analysis.py \
+--model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+--hf_dataset "lmms-lab/POPE" --dataset_split "test" \
+--dst-save-path ./excluded/pope_analysis.pkl --gpu-ids 0,1,2,3
+
+
+CUDA_VISIBLE_DEVICES=1,2,3 python analysis/gated_roi_analysis.py \
+    --model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+    --question-file /home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl \
+    --image-folder ../llava/playground/data/eval/textvqa/train_images \
+    --dst-save-path ./excluded/roi_KLdiv_textvqa_train_576res.pkl \
+    --gpu-ids 0,1,2
+    
+'''

QWENVL-PRIVATE/analysis/gen_roi_map_30b.py

@@ -0,0 +1,200 @@
+"""
+Generate ROI attention maps for Qwen3-VL-30B-A3B.
+Produces 10 GQA samples with grounding attention visualization.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python3 analysis/gen_roi_map_30b.py \
+        --model_path /mnt/bn/leonworkspace/terry/model/Qwen3-VL-30B-A3B-Instruct \
+        --num_samples 10
+
+Output: stage0/<model_name>/ with ROI map PNGs
+"""
+
+import os
+import sys
+import json
+import math
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+from tqdm import tqdm
+from transformers import AutoProcessor
+import argparse
+
+# Add project root to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from qwen_src.qwen3_vl import Qwen3VLForConditionalGeneration
+from qwen_src.mm_utils import expand2square
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map
+
+PATCH_SIZE = 32
+
+
+def plot_roi_map(image_path, grounding_attn, output_shape, question, answer, save_path):
+    """Plot image + ROI grounding attention map side by side."""
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4), dpi=200)
+
+    # Original image
+    if os.path.exists(image_path):
+        img = Image.open(image_path)
+        axes[0].imshow(img)
+    axes[0].set_title('Image', fontsize=10)
+    axes[0].axis('off')
+
+    # Grounding attention map
+    attn_2d = grounding_attn.reshape(output_shape)
+    im = axes[1].imshow(attn_2d, cmap='inferno', aspect='auto')
+    axes[1].set_title('ROI Grounding Attention', fontsize=10)
+    axes[1].axis('off')
+    plt.colorbar(im, ax=axes[1], shrink=0.6)
+
+    fig.suptitle(f'Q: {question[:80]}...\nA: {answer}', fontsize=8, y=0.02)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=200, bbox_inches='tight')
+    plt.close()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--num_samples", type=int, default=10)
+    parser.add_argument("--questions_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/llava_v1_5_mix665k_selected_qa.jsonl")
+    parser.add_argument("--images_root", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data")
+    args = parser.parse_args()
+
+    device = "cuda"
+    model_name = os.path.basename(args.model_path)
+    save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "stage0", model_name)
+    os.makedirs(save_dir, exist_ok=True)
+    print(f"Output dir: {save_dir}")
+
+    # Load model
+    print(f"Loading {model_name}...")
+    model = Qwen3VLForConditionalGeneration.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",  # need eager for attention weights
+        device_map="auto",  # auto device_map for 30B
+    ).eval()
+    print(f"Model loaded. Device map: {model.hf_device_map if hasattr(model, 'hf_device_map') else 'single GPU'}")
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        padding_side="left",
+        use_fast=True,
+    )
+
+    # Register hooks for attention extraction
+    activations, hooks = register_hooks(model, is_qwen3vl=True)
+
+    # Load GQA questions
+    print(f"Loading questions from {args.questions_path}...")
+    questions = [json.loads(l) for l in open(args.questions_path)]
+    gqa_questions = [q for q in questions if 'gqa' in q.get('image', '')]
+    print(f"Total: {len(questions)}, GQA: {len(gqa_questions)}")
+
+    # Model config
+    layers = model.config.text_config.num_hidden_layers
+    print(f"Model layers: {layers}")
+
+    # Process samples
+    for idx, line in enumerate(tqdm(gqa_questions[:args.num_samples], desc="Generating ROI maps")):
+        activations.clear()
+
+        image_file = line["image"]
+        question = line["text"] if "text" in line else line["prompt"]
+        image_path = os.path.join(args.images_root, image_file)
+
+        if not os.path.exists(image_path):
+            print(f"  Skip: {image_path} not found")
+            continue
+
+        # Load and prepare image
+        image = Image.open(image_path).convert('RGB')
+        processor.image_processor.max_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        processor.image_processor.min_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        image_size = image.size
+        image = expand2square(image, (127, 127, 127))
+
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": question},
+                {"type": "image", "image": image},
+            ]
+        }]
+
+        inputs = processor.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True,
+            return_dict=True, return_tensors="pt"
+        )
+        inputs = inputs.to(device)
+
+        # Generate with attention output
+        with torch.no_grad():
+            generated_ids = model.generate(**inputs, max_new_tokens=128, output_attentions=True)
+
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        answer = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+
+        # Build visual mask if not captured by hooks
+        if activations.get('visual_mask', None) is None:
+            vision_start_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+            vision_end_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+            input_ids_list = inputs['input_ids'].tolist()[0]
+            pos = input_ids_list.index(vision_start_id) + 1
+            pos_end = input_ids_list.index(vision_end_id)
+            activations['visual_mask'] = torch.zeros(
+                (1, inputs['input_ids'].shape[1]), dtype=torch.bool, device=device)
+            activations['visual_mask'][0, pos:pos_end] = True
+
+        output_shape = inputs['image_grid_thw'].cpu().numpy().squeeze(0)[1:] // 2
+        output_shape = tuple(output_shape.astype(int))
+
+        # Extract attention maps
+        attn_weights_map = get_attn_weights_map(activations, layers=layers, reduction='mean')
+
+        # Average grounding attention across middle-to-late layers
+        # For 30B (48 layers), use layers around 60-80% depth
+        grounding_layers = list(range(int(layers * 0.5), int(layers * 0.8)))
+        grounding_attn = np.zeros(output_shape)
+        count = 0
+        for layer_idx in grounding_layers:
+            key = f'layer{layer_idx}'
+            if key in attn_weights_map and 'output2images' in attn_weights_map[key]:
+                attn = attn_weights_map[key]['output2images']
+                if attn.size == np.prod(output_shape):
+                    grounding_attn += attn.reshape(output_shape)
+                    count += 1
+                else:
+                    # Multi-head: average across heads
+                    n_heads = attn.shape[0] // np.prod(output_shape)
+                    if n_heads > 0:
+                        attn_reshaped = attn.reshape(n_heads, *output_shape)
+                        grounding_attn += attn_reshaped.mean(axis=0)
+                        count += 1
+        if count > 0:
+            grounding_attn /= count
+
+        # Save visualization
+        save_path = os.path.join(save_dir, f"{idx}_roi_map.png")
+        plot_roi_map(image_path, grounding_attn, output_shape, question, answer, save_path)
+        print(f"  [{idx}] Q: {question[:60]}... A: {answer[:30]}...")
+
+    # Cleanup hooks
+    for h in hooks:
+        h.remove()
+
+    print(f"\nDone! {args.num_samples} ROI maps saved to {save_dir}/")
+
+
+if __name__ == "__main__":
+    main()

QWENVL-PRIVATE/analysis/get_o2i_attn.py

@@ -0,0 +1,107 @@
+from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
+import torch
+import matplotlib.pyplot as plt
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor
+import math
+
+def calculate_plt_size(attention_layer_num):
+    num_layers = attention_layer_num
+    cols = math.ceil(math.sqrt(num_layers))
+    rows = math.ceil(num_layers / cols)
+    return rows, cols
+
+device = 'cuda'
+model_path = 'Qwen/Qwen2.5-VL-3B-Instruct'
+
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model_path,
+                torch_dtype=torch.bfloat16,
+                attn_implementation="eager",
+            ).eval().to(device)
+processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True, padding_side='left', use_fast=True)
+
+image_path = './images/demo1.png'
+question = 'what is the date of the photo?'
+
+messages_query = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": image_path, "max_pixels": 512*28*28},
+            {"type": "text", "text": f"{question} Answer the question using a single word or phrase."},
+        ],
+    }
+]
+
+image_inputs, _ = process_vision_info(messages_query)
+
+text_query = processor.apply_chat_template(
+    messages_query,
+    tokenize=False,
+    add_generation_prompt=True
+)
+
+inputs = processor(
+    text=[text_query],
+    images=image_inputs,
+    padding=True,
+    return_tensors="pt",
+).to(device)
+
+messages_general = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image","image": image_path,"max_pixels": 512*28*28},
+            {"type": "text", "text": "Write a general description of the image. Answer the question using a single word or phrase."},
+        ],
+    }
+]
+
+text_general = processor.apply_chat_template(
+    messages_general,
+    tokenize=False,
+    add_generation_prompt=True
+)
+
+general_inputs = processor(
+    text=[text_general],
+    images=image_inputs,
+    padding=True,
+    return_tensors="pt",
+).to(device)
+
+image_inputs_aux = processor.image_processor(images=image_inputs)
+output_shape = image_inputs_aux["image_grid_thw"].numpy().squeeze(0)[1:]/2
+output_shape = output_shape.astype(int)
+
+with torch.no_grad():
+    vision_start_token_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+    vision_end_token_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+    pos = inputs['input_ids'].tolist()[0].index(vision_start_token_id) + 1
+    pos_end = inputs['input_ids'].tolist()[0].index(vision_end_token_id)
+
+    output = model(**inputs, output_attentions=True)
+    general_output = model(**general_inputs, output_attentions=True)
+    rows, cols = calculate_plt_size(len(output.attentions))
+    fig, axes = plt.subplots(rows, cols, figsize=(10.8, 16))
+    for i, ax in enumerate(axes.flatten()):
+        if i < len(output.attentions):
+            att = output.attentions[i][0, :, -1, pos:pos_end].mean(dim=0)
+            att = att.to(torch.float32).detach().cpu().numpy()
+
+            general_att = general_output.attentions[i][0, :, -1, pos:pos_end].mean(dim=0)
+            general_att = general_att.to(torch.float32).detach().cpu().numpy()
+
+            att = att / general_att
+
+            ax.imshow(att.reshape(output_shape), cmap="viridis", interpolation="nearest")
+            ax.set_title(f"Layer {i+1}")
+            ax.axis("off")
+        else:
+            ax.axis("off")
+
+    plt.tight_layout()
+    plt.show()
+    pass

QWENVL-PRIVATE/analysis/llm-as-judge-roi-samples.py

@@ -0,0 +1,383 @@
+import os
+import json
+import pickle
+import argparse
+import math
+import torch
+import multiprocessing as mp
+import re
+from torch.utils.data import Dataset, DataLoader
+# CHANGED: Import standard LLM classes instead of VL classes
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from tqdm import tqdm
+import glob
+
+
+# --- 1. Dataset & Collator ---
+
+class JudgeDataset(Dataset):
+    def __init__(self, items):
+        self.items = items
+
+    def __len__(self):
+        return len(self.items)
+
+    def __getitem__(self, idx):
+        item = self.items[idx]
+
+        # 1. Format Ground Truth
+        gt = item['ground_truth']
+        if isinstance(gt, list):
+            gt_str = ", ".join([str(x) for x in gt])
+        else:
+            gt_str = str(gt)
+
+        # 2. Judge Prompt
+        prompt_text = (
+            f"You are an intelligent evaluator for a Visual Question Answering task.\n"
+            f"Compare the Model Predictions against the Ground Truth.\n\n"
+
+            f"--- DATA ---\n"
+            f"Question: {item['question']}\n"
+            f"Ground Truth: {gt_str}\n"
+            f"Model A Prediction: {item['base_resp']}\n"
+            f"Model B Prediction: {item['roi_resp']}\n\n"
+
+            f"--- CRITERIA ---\n"
+            f"1. **Containment**: If the Ground Truth is a short value (e.g., '6000') and the Model Prediction contains it correctly within a sentence (e.g., 'The value is 6000'), mark it as **Yes**.\n"
+            f"2. **Synonyms**: Accept standard synonyms (e.g., 'Yes' = 'True', '10%' = '0.1').\n"
+            f"3. **Formatting**: Ignore punctuation or capitalization differences (e.g., 'april' = 'April').\n"
+            f"4. **Contradiction**: Only mark 'No' if the prediction is factually wrong or missing the key information.\n\n"
+
+            f"--- TASK ---\n"
+            f"Does Model A provide the correct answer? (Yes/No)\n"
+            f"Does Model B provide the correct answer? (Yes/No)\n\n"
+
+            f"Respond in this exact format:\n"
+            f"A: [Yes/No]\n"
+            f"B: [Yes/No]"
+        )
+
+        messages = [
+            {"role": "system",
+             "content": "You are a fair and robust evaluator. You focus on semantic correctness over strict string matching."},
+            {"role": "user", "content": prompt_text}
+        ]
+
+        return {
+            "messages": messages,
+            "original_item": item
+        }
+
+
+class JudgeCollator:
+    def __init__(self, model_path):
+        self.model_path = model_path
+        self.tokenizer = None
+
+    def __call__(self, batch):
+        # CHANGED: Initialize AutoTokenizer for pure LLM
+        if self.tokenizer is None:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_path,
+                trust_remote_code=True,
+                padding_side='left',
+                use_fast=True
+            )
+            # Ensure pad token exists for batched inference
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        raw_messages = [item['messages'] for item in batch]
+        original_items = [item['original_item'] for item in batch]
+
+        # CHANGED: Apply chat template for text only
+        text_inputs = [
+            self.tokenizer.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+            for msg in raw_messages
+        ]
+
+        inputs = self.tokenizer(
+            text_inputs,
+            padding=True,
+            return_tensors="pt"
+        )
+
+        return {
+            "inputs": inputs,
+            "items": original_items
+        }
+
+
+# --- 2. Worker Functions ---
+
+def worker_process(chunk, gpu_id, args):
+    device = f'cuda:{gpu_id}'
+
+    # Storage for results
+    regressions_found = []  # Base: Right, ROI: Wrong
+    improvements_found = []  # Base: Wrong, ROI: Right
+
+    # Setup Dataset & Loader
+    dataset = JudgeDataset(chunk)
+    collator = JudgeCollator(args.model_path)
+
+    dataloader = DataLoader(
+        dataset,
+        # NOTE: 32B model requires more VRAM. You might need to lower batch_size manually if OOM occurs.
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        collate_fn=collator,
+        pin_memory=True
+    )
+
+    if len(dataset) == 0:
+        return {"regressions": [], "improvements": []}
+
+    try:
+        # Load Model
+        print(f"Worker {gpu_id}: Loading Judge Model (LLM)...")
+        # CHANGED: Use AutoModelForCausalLM
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device
+        )
+        model.eval()
+
+        # Access tokenizer from collator
+        tokenizer = collator.tokenizer if collator.tokenizer else None
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
+
+        print(f"Worker {gpu_id}: Starting Judging...")
+
+        for batch_data in tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id):
+            inputs = batch_data['inputs']
+            items = batch_data['items']
+
+            # Move inputs to device (AutoTokenizer outputs input_ids and attention_mask)
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                output_ids = model.generate(
+                    **inputs,
+                    max_new_tokens=32,
+                    do_sample=False,
+                    pad_token_id=tokenizer.eos_token_id
+                )
+
+            # CHANGED: Slice outputs correctly for LLM (remove prompt tokens)
+            input_len = inputs['input_ids'].shape[1]
+            generated_ids = output_ids[:, input_len:]
+            responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+
+            for item, resp in zip(items, responses):
+                clean_resp = resp.replace("\n", " ").strip()
+
+                # Parse A (Baseline)
+                base_match = re.search(r"A\s*[:\-\.]\s*(Yes|True|Correct)", clean_resp, re.IGNORECASE)
+                base_correct = bool(base_match)
+
+                # Parse B (ROI)
+                roi_match = re.search(r"B\s*[:\-\.]\s*(Yes|True|Correct)", clean_resp, re.IGNORECASE)
+                roi_correct = bool(roi_match)
+
+                # Helper to create result dict
+                def create_result_entry():
+                    s_bbox = item.get('surrounding_bbox')
+                    if hasattr(s_bbox, 'tolist'): s_bbox = s_bbox.tolist()
+
+                    r_mask = item.get('roi_mask')
+                    mask_shape = r_mask.shape if hasattr(r_mask, 'shape') else "N/A"
+
+                    return {
+                        "dataset": item.get('dataset', 'unknown'),
+                        "line_number": item['line_number'],
+                        "question_id": item['id'],
+                        "image_id": item['image'],
+                        "question": item['question'],
+                        "ground_truth": item['ground_truth'],
+                        "baseline_response": item['base_resp'],
+                        "roi_response": item['roi_resp'],
+                        "judge_raw_output": resp,
+                        "surrounding_bbox": s_bbox,
+                        "roi_mask_shape": mask_shape
+                    }
+
+                # CASE 1: REGRESSION (Base Right, ROI Wrong)
+                if base_correct and not roi_correct:
+                    regressions_found.append(create_result_entry())
+
+                # CASE 2: IMPROVEMENT (Base Wrong, ROI Right)
+                elif not base_correct and roi_correct:
+                    improvements_found.append(create_result_entry())
+
+            del inputs, output_ids, generated_ids
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: Critical Error: {e}")
+        import traceback
+        traceback.print_exc()
+
+    return {"regressions": regressions_found, "improvements": improvements_found}
+
+
+def worker_wrapper(chunk, gpu_id, args, return_dict):
+    result = worker_process(chunk, gpu_id, args)
+    return_dict[gpu_id] = result
+
+
+# --- 3. Main Execution ---
+
+def load_all_paired_data(search_dir):
+    all_base_data = {}
+    all_roi_data = {}
+
+    search_pattern = os.path.join(search_dir, "*_qwen2.5vl-7B.pkl")
+    base_files = glob.glob(search_pattern)
+
+    print(f"Found {len(base_files)} baseline files in '{search_dir}'. Loading pairs...")
+
+    for base_file in base_files:
+        filename = os.path.basename(base_file)
+        dataset_name = filename.replace("_qwen2.5vl-7B.pkl", "")
+        roi_file = base_file.replace(".pkl", "-roi.pkl")
+
+        if not os.path.exists(roi_file):
+            print(f"⚠️  Skipping {dataset_name}: ROI file not found ({roi_file})")
+            continue
+
+        try:
+            with open(base_file, 'rb') as f:
+                cur_base = pickle.load(f)
+            with open(roi_file, 'rb') as f:
+                cur_roi = pickle.load(f)
+
+            for k in cur_base.keys():
+                if k in cur_roi:
+                    unique_key = f"{dataset_name}_{k}"
+                    all_base_data[unique_key] = cur_base[k]
+                    all_roi_data[unique_key] = cur_roi[k]
+                    all_base_data[unique_key]['dataset'] = dataset_name
+                    all_roi_data[unique_key]['dataset'] = dataset_name
+
+            print(f"  + Loaded {dataset_name}: {len(cur_base)} samples")
+
+        except Exception as e:
+            print(f"❌ Error loading {dataset_name}: {e}")
+
+    return all_base_data, all_roi_data
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    # CHANGED: Default model path now points to a pure LLM path example
+    parser.add_argument('--model-path', type=str, default='/root/public/Qwen/Qwen2.5-32B-Instruct')
+    parser.add_argument('--data-dir', type=str, default='./', help="Directory containing .pkl files")
+    parser.add_argument('--save-path', type=str, default="regressions.json",
+                        help="File for Regressions (Base Right, ROI Wrong)")
+    parser.add_argument('--save-path-improvements', type=str, default="improvements.json",
+                        help="File for Improvements (Base Wrong, ROI Right)")
+    parser.add_argument('--gpu-ids', type=str, default='0,1,2,3')
+    parser.add_argument('--batch-size', type=int, default=8)
+    parser.add_argument('--num_workers', type=int, default=4)
+    args = parser.parse_args()
+
+    # 1. Load All Data Pairs
+    base_data, roi_data = load_all_paired_data(args.data_dir)
+
+    if not base_data:
+        print("No paired data found. Exiting.")
+        return
+
+    # 2. Prepare Items
+    judge_items = []
+    print(f"Preparing {len(base_data)} items for judging...")
+
+    for k, base_item in base_data.items():
+        roi_item = roi_data[k]
+
+        # Optimization: Skip identical strings (Result is definitely same, so neither reg nor imp)
+        if str(base_item.get('response', '')).strip() == str(roi_item.get('response', '')).strip():
+            continue
+
+        judge_items.append({
+            'unique_key': k,
+            'dataset': base_item.get('dataset', 'unknown'),
+            'line_number': base_item.get('line_number', 'unknown'),
+            'id': base_item['id'],
+            'image': base_item['image'],
+            'question': base_item['prompt'],
+            'ground_truth': base_item['text'],
+            'base_resp': base_item.get('response', ''),
+            'roi_resp': roi_item.get('response', ''),
+            'surrounding_bbox': roi_item.get('surrounding_bbox'),
+            'roi_mask': roi_item.get('roi_mask')
+        })
+
+    print(f"Items requiring judgement (non-identical): {len(judge_items)}")
+
+    # 3. Parallel Processing
+    try:
+        mp.set_start_method('spawn', force=True)
+    except RuntimeError:
+        pass
+
+    gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
+    num_gpus = len(gpu_ids)
+
+    chunk_size = math.ceil(len(judge_items) / num_gpus)
+    chunks = [judge_items[i:i + chunk_size] for i in range(0, len(judge_items), chunk_size)]
+
+    manager = mp.Manager()
+    return_dict = manager.dict()
+    processes = []
+
+    print(f"Starting judging on {num_gpus} GPUs...")
+    for i in range(num_gpus):
+        chunk_data = chunks[i] if i < len(chunks) else []
+        p = mp.Process(target=worker_wrapper, args=(chunk_data, gpu_ids[i], args, return_dict))
+        p.start()
+        processes.append(p)
+
+    for p in processes:
+        p.join()
+
+    # 4. Aggregation
+    all_regressions = []
+    all_improvements = []
+
+    for gpu_id, res in return_dict.items():
+        all_regressions.extend(res['regressions'])
+        all_improvements.extend(res['improvements'])
+
+    print(f"\n✅ Analysis Complete.")
+    print(f"  - Regressions (Worse): {len(all_regressions)}")
+    print(f"  - Improvements (Better): {len(all_improvements)}")
+
+    # Save Regressions
+    with open(args.save_path, 'w', encoding='utf-8') as f:
+        json.dump(all_regressions, f, indent=2)
+    print(f"Saved regressions to {args.save_path}")
+
+    # Save Improvements
+    with open(args.save_path_improvements, 'w', encoding='utf-8') as f:
+        json.dump(all_improvements, f, indent=2)
+    print(f"Saved improvements to {args.save_path_improvements}")
+
+
+if __name__ == "__main__":
+    main()
+
+'''
+python analysis/llm-as-judge-roi-samples.py \
+    --model-path /root/public/Qwen/Qwen2.5-32B-Instruct \
+    --data-dir /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining  \
+    --save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/regressions-v2-32b-judge.json \
+    --save-path-improvements /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/improvements-v2-32b-judge.json \
+    --batch-size 8 \
+    --gpu-ids 0,1,2,3
+'''

QWENVL-PRIVATE/analysis/making_base_model_response.py

@@ -0,0 +1,544 @@
+import sys
+import os
+script_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(script_dir)
+sys.path.append(project_root)
+
+from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
+from qwen_src.qwen2_vl import Qwen2VLForConditionalGeneration
+try:
+    from qwen_src.qwen3_vl.modeling_qwen3_vl_batch import Qwen3VLForConditionalGeneration
+except Exception:
+    Qwen3VLForConditionalGeneration = None
+import torch
+import matplotlib.pyplot as plt
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor
+import math
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+
+import json
+import os
+from tqdm import tqdm
+
+import pickle
+import argparse
+import multiprocessing as mp
+import pandas as pd  # <-- Make sure to import pandas
+from matching_utils import TextVQAAccuracyEvaluator
+from check_gate_label import plot_single_sample
+
+
+def infer_model_type(model_path, model_type):
+    if model_type and model_type != "auto":
+        return model_type
+    lower = model_path.lower()
+    if "qwen3" in lower:
+        return "qwen3"
+    if "qwen2.5" in lower or "qwen2_5" in lower or "qwen2-5" in lower:
+        return "qwen2_5"
+    if "qwen2" in lower:
+        return "qwen2"
+    return "qwen2_5"
+
+
+def resolve_processor_path(model_path, model_type, processor_path=None):
+    if processor_path:
+        return processor_path
+    lower = model_path.lower()
+    public_prefix = "/root/public/Qwen"
+
+    def to_repo_path(name):
+        if "/root/public" in model_path:
+            return f"{public_prefix}/{name}"
+        return f"Qwen/{name}"
+
+    if model_type == "qwen3":
+        if "8b" in lower:
+            return to_repo_path("Qwen3-VL-8B-Instruct")
+        if "4b" in lower:
+            return to_repo_path("Qwen3-VL-4B-Instruct")
+    if model_type == "qwen2_5":
+        if "7b" in lower:
+            return to_repo_path("Qwen2.5-VL-7B-Instruct")
+        if "3b" in lower:
+            return to_repo_path("Qwen2.5-VL-3B-Instruct")
+    if model_type == "qwen2" and "7b" in lower:
+        return to_repo_path("Qwen2-VL-7B-Instruct")
+
+    return model_path
+
+
+def resolve_model_class(model_type):
+    if model_type == "qwen3":
+        if Qwen3VLForConditionalGeneration is None:
+            raise RuntimeError("Qwen3VLForConditionalGeneration not available.")
+        return Qwen3VLForConditionalGeneration
+    if model_type == "qwen2":
+        return Qwen2VLForConditionalGeneration
+    return Qwen2_5_VLForConditionalGeneration
+
+
+def set_twig_flags(model, model_type, two_stage_flag):
+    if model_type == "qwen3":
+        if hasattr(model, "enable_twig"):
+            model.enable_twig = two_stage_flag
+        if hasattr(model, "roi_enable2stage"):
+            model.roi_enable2stage = two_stage_flag
+        if hasattr(model, "language_model"):
+            if hasattr(model.language_model, "enable_twig"):
+                model.language_model.enable_twig = two_stage_flag
+            if hasattr(model.language_model, "roi_enable2stage"):
+                model.language_model.roi_enable2stage = two_stage_flag
+        return
+
+    if hasattr(model, "model"):
+        if hasattr(model.model, "enable_twig"):
+            model.model.enable_twig = two_stage_flag
+        if hasattr(model.model, "roi_enable2stage"):
+            model.model.roi_enable2stage = two_stage_flag
+        if hasattr(model.model, "two_stage_vanilla"):
+            model.model.two_stage_vanilla = two_stage_flag
+    if hasattr(model, "roi_enable2stage"):
+        model.roi_enable2stage = two_stage_flag
+def get_KL_divergence(logits_list,):
+    import torch.nn as nn
+    from torch.nn import functional as F
+    KL_div = nn.KLDivLoss(reduction='batchmean')
+    KL_div_list = [KL_div(F.log_softmax(logits_list[-1], dim=1), F.softmax(logits_list[i], dim=1)).item() for i in range(len(logits_list))]
+    return KL_div_list
+
+
+from torch.utils.data import Dataset, DataLoader
+
+
+def worker_wrapper(chunk, gpu_id, args, return_dict):
+    """
+    Wraps the main worker function to handle return values via shared dictionary.
+    """
+
+    result = worker_process(chunk, gpu_id, args)
+    # Store result in the shared dictionary using gpu_id as key
+    return_dict[gpu_id] = result
+
+
+
+class QwenEvalDataset(Dataset):
+    def __init__(self, questions_chunk, args):
+        self.data = questions_chunk
+        self.args = args
+        self.image_folder = args.image_folder
+        self.use_short_prompt = getattr(args, 'use_short_prompt', False)
+
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        line = self.data[idx]
+        image_file_str = ''
+        ground_truth = None
+        qs_text = ""
+        image = None
+
+        try:
+
+            qs_text = line["text"]
+            qs_text = qs_text
+            image_file = line["image"]
+            image_path = os.path.join(self.image_folder, image_file)
+            image = Image.open(image_path).convert('RGB')
+            image_file_str = image_file
+            ground_truth = line.get("answer") if "answer" in line else line.get("answers", None)
+
+            if "<image>" in qs_text:
+                qs_text = qs_text.replace("<image>", "")  # Remove it from the text body
+
+            qs = "<image>\n" + qs_text
+
+            if self.use_short_prompt:
+                qs = qs + "\nAnswer the question using a single word or phrase."
+
+            # Return a dict containing everything needed for inference
+            return {
+                'image': image,
+                'prompt_text': qs,
+                'meta': {
+                    'image_id': image_file_str,
+                    'prompt': qs_text,
+                    'ground_truth': ground_truth,
+                    'line_number': line['original_line_number'],
+                    'id': line.get('question_id', line.get('original_line_number', 'unknown'))
+                },
+                'valid': True
+            }
+
+        except Exception as e:
+            print(f"Error loading item {idx}: {e}")
+            # Return dummy invalid item
+            return {'valid': False, 'meta': {'line_number': line.get('original_line_number', -1)}}
+
+
+class SingleResCollator:
+    def __init__(self, processor_path, patch_size):
+        self.processor_path = processor_path
+        self.processor = None
+        # Define the requested limits
+        self.min_pixels = 128 * patch_size * patch_size
+        self.max_pixels = 576 * patch_size * patch_size
+
+    def __call__(self, batch):
+        if self.processor is None:
+            # Initialize processor once per worker
+            self.processor = AutoProcessor.from_pretrained(
+                self.processor_path,
+                trust_remote_code=True,
+                padding_side='left',
+                use_fast=True
+            )
+            # Set the resolution constraints globally for this processor
+            self.processor.image_processor.min_pixels = self.min_pixels
+            self.processor.image_processor.max_pixels = self.max_pixels
+
+        batch = [item for item in batch if item.get('valid', False)]
+        if not batch:
+            return None
+
+        # 1. Extract raw data
+        raw_images = [item['image'] for item in batch]
+        raw_prompts = [item['prompt_text'] for item in batch]
+        metas = [item['meta'] for item in batch]
+
+        # 2. Prepare text and images
+        current_texts = []
+        for prompt_str, img in zip(raw_prompts, raw_images):
+            # Qwen-VL-Chat template construction
+            messages = [{"role": "user", "content": [{"type": "text", "text": prompt_str}, {"image": img}]}]
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+            # Specific replacement logic for Qwen2.5-VL
+            text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '') \
+                .replace('<image>', '<image><|vision_start|><|image_pad|><|vision_end|>')
+            current_texts.append(text)
+
+        try:
+            # 3. Process Batch (Single Resolution Setting)
+            inputs = self.processor(
+                text=current_texts,
+                images=raw_images,
+                padding=True,
+                return_tensors="pt"
+            )
+            # Attach extra data needed later if necessary (or just return inputs)
+            inputs.data['src_images'] = raw_images
+
+            return {
+                'metas': metas,
+                'inputs': inputs
+            }
+
+        except Exception as e:
+            print(f"Error processing batch: {e}")
+            return None
+
+
+def worker_process(questions_chunk, gpu_id, args):
+    device = f'cuda:{gpu_id}'
+    answer_analysis_dict = {}  # Renamed from ms_answer_analysis_dict
+    max_new_tokens = 256
+
+    # --- SETUP DATALOADER ---
+    batch_size = getattr(args, 'batch_size', 8)
+    num_workers = getattr(args, 'num_workers', 4)
+
+    dataset = QwenEvalDataset(questions_chunk, args)
+
+    # Initialize the NEW single resolution collator
+    collator = SingleResCollator(args.processor_path, args.patch_size)
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collator,
+        pin_memory=True
+    )
+
+    if len(dataset) == 0: return answer_analysis_dict
+
+    try:
+        # Load Model
+        print(f"Worker {gpu_id}: Loading model...")
+        model_class = resolve_model_class(args.model_type)
+        model = model_class.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device
+        )
+        model.eval()
+
+        # Flags (keeping your original config)
+        two_stage_flag = False
+        set_twig_flags(model, args.model_type, two_stage_flag)
+
+        # Tokenizer for decoding
+        tokenizer = AutoProcessor.from_pretrained(args.processor_path, trust_remote_code=True).tokenizer
+
+        print(f"Worker {gpu_id}: Starting inference...")
+
+        # --- MAIN INFERENCE LOOP ---
+        progress_bar = tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id)
+
+        for batch_data in progress_bar:
+            if batch_data is None: continue
+
+            metas = batch_data['metas']
+            inputs = batch_data['inputs']  # Now getting direct inputs, not inputs_by_res
+
+            try:
+                # Move to GPU
+                inputs = inputs.to(device)
+
+                with torch.no_grad():
+                    output_ids = model.generate(
+                        **inputs,
+                        max_new_tokens=max_new_tokens,
+                        output_attentions=False,
+                        pad_token_id=tokenizer.eos_token_id
+                    )
+
+                # Decode
+                input_ids_len = inputs.input_ids.shape[1]
+                generated_ids = output_ids[:, input_ids_len:]
+                batch_answers = tokenizer.batch_decode(
+                    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+                )
+
+                # --- SAVE RESULTS ---
+                for idx, (meta, answer) in enumerate(zip(metas, batch_answers)):
+                    line_number = meta['line_number']
+
+                    tmp_sample = {
+                        'image': meta['image_id'],
+                        'prompt': meta['prompt'],
+                        'text': meta['ground_truth'],
+                        'id': meta['id'],
+                        'response': answer  # Storing single response
+                    }
+                    answer_analysis_dict[line_number] = tmp_sample
+
+                # Cleanup GPU memory for this batch
+                del inputs, output_ids, generated_ids
+
+            except RuntimeError as e:
+                msg = "ERROR_OOM" if "out of memory" in str(e) else f"ERROR_RUNTIME: {e}"
+                print(f"Worker {gpu_id}: {msg}")
+                torch.cuda.empty_cache()
+                # Log errors for this batch
+                for meta in metas:
+                    line_number = meta['line_number']
+                    answer_analysis_dict[line_number] = {
+                        'id': meta['id'],
+                        'response': msg
+                    }
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: Critical Error: {e}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        torch.cuda.empty_cache()
+
+    return answer_analysis_dict
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Parallel KL Divergence Evaluation")
+    parser.add_argument('--model-path', type=str, required=True,
+                        help="Path to the model and processor")
+    parser.add_argument('--model-type', type=str, default="auto",
+                        choices=["auto", "qwen2_5", "qwen2", "qwen3"],
+                        help="Model family for loading (auto detects from model path).")
+    parser.add_argument('--processor-path', type=str, default="",
+                        help="Override processor path if auto detection is incorrect.")
+    parser.add_argument('--question-file', type=str, default=None,
+                        help="Path to the JSONL question file (E.g., TextVQA)")
+    parser.add_argument('--image-folder', type=str, default=None,
+                        help="Path to the image directory (for JSONL file)")
+    parser.add_argument('--dst-save-path', type=str, required=True,
+                        help="Path to save the output .pkl file")
+    parser.add_argument('--gpu-ids', type=str, default='0,1',
+                        help="Comma-separated list of GPU IDs to use (e.g., '0,1,2,3')")
+    # --- ADDED: New HF Dataset args ---
+    parser.add_argument('--hf_dataset', type=str, default=None,
+                        help="Name of the Hugging Face dataset (e.g., 'craigwu/vstar_bench')")
+    parser.add_argument('--dataset_split', type=str, default='test',
+                        help="Split to use from the Hugging Face dataset (e.g., 'test')")
+    parser.add_argument('--is_check_pred', type=bool, default=False,
+                        help="Whether to check high-res predictions")
+    ### --- END ADDITION ---
+    parser.add_argument('--use-short-prompt', action='store_true',
+                        help="Whether to use short prompt format")
+    args = parser.parse_args()
+    args.model_type = infer_model_type(args.model_path, args.model_type)
+    args.processor_path = resolve_processor_path(
+        args.model_path, args.model_type, args.processor_path
+    )
+    args.patch_size = 32 if args.model_type == "qwen3" else 28
+
+
+    # --- 1. Define a Wrapper Function ---
+    # This helper must be defined at the top level (global scope) so it can be pickled.
+    # --- 2. The Main Execution Block ---
+    # Ensure start method is spawn (Crucial for PyTorch+CUDA)
+    try:
+        mp.set_start_method('spawn', force=True)
+    except RuntimeError:
+        pass
+
+    # 1. Prepare Data and Split into Chunks
+    gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
+    num_gpus = len(gpu_ids)
+    questions = []
+
+    if args.hf_dataset:
+        print(f"Loading Hugging Face dataset {args.hf_dataset}, split {args.dataset_split}")
+        from datasets import load_dataset
+
+        dataset = load_dataset(args.hf_dataset, split=args.dataset_split)
+        print("Casting 'image' column to load from cache...")
+        questions_list = list(dataset)
+
+        new_questions = []
+        for i, tmp_ques in enumerate(questions_list):
+            tmp_ques['original_line_number'] = i
+            new_questions.append(tmp_ques)
+        questions = new_questions
+        print(f"Loaded {len(questions)} preprocessed questions from Hugging Face.")
+
+    elif args.question_file:
+        print("Loading and preprocessing questions from JSONL file...")
+        if not args.image_folder:
+            raise ValueError("--image-folder is required when using --question-file")
+
+        questions_list = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+
+        new_questions = []
+        for i, tmp_ques in enumerate(questions_list):
+            qs = tmp_ques['text'] if 'text' in tmp_ques else tmp_ques['prompt']
+            if 'Reference OCR' in qs:
+                qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
+            tmp_ques['text'] = qs
+            tmp_ques['original_line_number'] = i
+            new_questions.append(tmp_ques)
+        questions = new_questions
+        print(f"Loaded {len(questions)} preprocessed questions from local file.")
+
+    else:
+        raise ValueError("You must provide either --question-file or --hf_dataset")
+
+    # Create chunks
+    chunk_size = math.ceil(len(questions) / num_gpus)
+    chunks = [questions[i:i + chunk_size] for i in range(0, len(questions), chunk_size)]
+
+    # 2. Run processes in parallel using mp.Process (Manual Management)
+    print(f"Starting parallel processing on {num_gpus} GPUs: {gpu_ids}")
+
+    # Create a Manager to hold the results safely across processes
+    manager = mp.Manager()
+    return_dict = manager.dict()
+    processes = []
+
+    for i in range(num_gpus):
+        chunk_data = chunks[i] if i < len(chunks) else []
+
+        # We pass the shared 'return_dict' to the wrapper
+        p = mp.Process(
+            target=worker_wrapper,
+            args=(chunk_data, gpu_ids[i], args, return_dict)
+        )
+        p.start()
+        processes.append(p)
+
+    # Wait for all processes to finish
+    for p in processes:
+        p.join()
+
+    # 3. Gather and Save Results
+    print("All chunks processed. Gathering results...")
+    final_dict = {}
+
+    # Extract results from the shared return_dict
+    for gpu_id, sub_dict in return_dict.items():
+        final_dict.update(sub_dict)
+
+    #check save directory
+    os.makedirs(os.path.dirname(os.path.expanduser(args.dst_save_path)), exist_ok=True)
+    with open(os.path.expanduser(args.dst_save_path), 'wb') as f:
+        pickle.dump(final_dict, f)
+
+    print(f"✅ Successfully processed and saved {len(final_dict)} answers to {args.dst_save_path}.")
+
+
+'''
+CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
+    --model-path Qwen/Qwen2.5-VL-7B-Instruct \
+    --question-file /home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl \
+    --image-folder ../llava/playground/data/eval/textvqa/train_images \
+    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B.pkl \
+    --gpu-ids 0,1
+
+CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
+    --model-path Qwen/Qwen2.5-VL-7B-Instruct \
+    --question-file /home/yuheng/datasets/vstar_coco_spatial_relation_data.jsonl \
+    --image-folder /home/yuheng/datasets \
+    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B.pkl \
+    --gpu-ids 0,1
+
+
+
+
+python analysis/making_base_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/textvqa/converted_llava_style_train.jsonl \
+    --image-folder /root/epfs/sproject/datasets/textvqa/train_images \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+    
+python analysis/making_base_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/chartqa_28k_qa.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/chartqa_qwen2.5vl-7B-short-prompt.pkl \\
+    --use-short-prompt \
+    --gpu-ids 0,1
+
+python analysis/making_base_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/visual_cot_infovqa_subset15k.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_infovqa_qwen2.5vl-7B-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+    
+python analysis/making_base_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/visual_cot_docvqa_subset33k.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_docvqa_qwen2.5vl-7B-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+    
+python analysis/making_base_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/vstar_coco_spatial_relation_data.jsonl \
+    --image-folder /root/epfs/sproject/datasets/ \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+'''

QWENVL-PRIVATE/analysis/making_gated_training_samples.py

@@ -0,0 +1,781 @@
+import sys
+import os
+script_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(script_dir)
+sys.path.append(project_root)
+
+from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
+from qwen_src.qwen2_vl import Qwen2VLForConditionalGeneration
+import torch
+import matplotlib.pyplot as plt
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor
+import math
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+
+import json
+import os
+from tqdm import tqdm
+
+import pickle
+import argparse
+import multiprocessing as mp
+import pandas as pd  # <-- Make sure to import pandas
+from matching_utils import TextVQAAccuracyEvaluator
+from check_gate_label import plot_single_sample
+def get_KL_divergence(logits_list,):
+    import torch.nn as nn
+    from torch.nn import functional as F
+    KL_div = nn.KLDivLoss(reduction='batchmean')
+    KL_div_list = [KL_div(F.log_softmax(logits_list[-1], dim=1), F.softmax(logits_list[i], dim=1)).item() for i in range(len(logits_list))]
+    return KL_div_list
+
+
+from torch.utils.data import Dataset, DataLoader
+
+
+def worker_wrapper(chunk, gpu_id, args, return_dict):
+    """
+    Wraps the main worker function to handle return values via shared dictionary.
+    """
+    if args.is_check_pred:
+        result = check_high_res_prediction(chunk, gpu_id, args)
+    else:
+        result = worker_process_ms(chunk, gpu_id, args)
+    # Store result in the shared dictionary using gpu_id as key
+    return_dict[gpu_id] = result
+
+
+def worker_process_kl_divergence(questions_chunk, gpu_id, args):
+    """
+    This is the worker function that runs on a single GPU.
+    It loads the model and processes its assigned chunk of questions.
+    """
+
+    device = f'cuda:{gpu_id}'
+    KL_div_analysis_dict = {}
+    max_new_tokens = 1024
+    if not questions_chunk:
+        print(f"Worker {gpu_id}: No questions in chunk. Exiting.")
+        return KL_div_analysis_dict
+
+    try:
+        # 1. Load Model and Processor INSIDE the worker
+        print(f"Worker {gpu_id}: Loading processor from {args.model_path}...")
+        processor = AutoProcessor.from_pretrained(
+            'Qwen/Qwen2.5-VL-7B-Instruct',
+            trust_remote_code=True,
+            padding_side='left', use_fast=True
+        )
+
+        print(f"Worker {gpu_id}: Loading model from {args.model_path}...")
+        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device  # Assign model directly to the correct device
+        )
+        model.eval()
+
+        # Set processor image size limits (matching your original function)
+        processor.image_processor.max_pixels = 576 * 28 * 28
+        processor.image_processor.min_pixels = 576 * 28 * 28
+
+        # 2. Process the chunk of questions
+        print(f"Worker {gpu_id}: Starting processing of {len(questions_chunk)} items.")
+
+        # Use tqdm for progress, disable for all but the first worker if too noisy
+        progress_bar = tqdm(questions_chunk, desc=f"GPU {gpu_id}", position=gpu_id)
+        base_answer = two_stage_answer = ""  # Initialize to avoid reference before assignment
+        for line in progress_bar:
+            try:
+                # --- MODIFIED: Handle different dataset formats ---
+                image_file_str = ''
+                ground_truth = None
+                if args.question_file is not None and 'textvqa' in args.question_file:  # Case 1: Old format with "text" key
+                    qs_text = line["text"]
+                    qs_text = qs_text + " Answer the question using a single word or phrase." if "Answer the question using a single word or phrase." not in qs_text else qs_text
+                    image_file = line["image"]
+                    image_path = os.path.join(args.image_folder, image_file)
+                    image = Image.open(image_path).convert('RGB')
+                    image_file_str = image_file
+                    ground_truth = line.get("answers")  # Store the ground truth
+                    assert ground_truth is not None, "Ground truth answers missing in the data."
+                elif args.question_file is not None and 'visual_cot' in args.question_file:  # Case 1.5: Old format with "question" key
+                    qs_text = line["text"]
+                    qs_text = qs_text + " Answer the question using a single word or phrase." if "Answer the question using a single word or phrase." not in qs_text else qs_text
+                    image_file = line["image"]
+                    image_path = os.path.join(args.image_folder, image_file)
+                    image = Image.open(image_path).convert('RGB')
+                    image_file_str = image_file
+                    ground_truth = line.get("answer")
+                elif args.hf_dataset is not None:  # Case 2: New HF Dataset format (e.g., vstar_bench)
+                    if 'vstar' in args.hf_dataset:
+                        qs_text = line["text"]
+                        image = line["image"].convert('RGB')
+                        ground_truth = line.get("label")  # Store the ground truth
+                        image_file_str = f"hf_dataset_{line.get('question_id', 'unknown')}"
+                    elif 'POPE' in args.hf_dataset:
+                        qs_text = line["question"] + " Answer the question using a single word or phrase."
+                        image = line["image"].convert('RGB')
+                        ground_truth = line.get("answer")
+                        image_file_str = f"hf_dataset_{line.get('question_id', 'unknown')}"
+
+                else:
+                    print(f"Worker {gpu_id}: Skipping line {line.get('original_line_number')}, unknown format.")
+                    continue
+                # --- END MODIFICATION ---
+                qs = "<image>\n" + qs_text
+                messages = [
+                    {"role": "user", "content": [{"type": "text", "text": f"{qs}"}, {"image": image}]}
+                ]
+                text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+                text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '').replace('<image>',
+                                                                                               '<image><|vision_start|><|image_pad|><|vision_end|>')
+
+                inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt").to(device)
+
+                # We don't need these for KL div, but keeping if your model.forward needs them
+                inputs.data['src_images'] = [image]
+                inputs.data['processor'] = processor
+                input_ids_len = inputs.input_ids.shape[1]
+
+                # --- Run with two_stage_flag = False (Base) ---
+                two_stage_flag = False
+                model.model.enable_twig = two_stage_flag
+                model.model.roi_enable2stage = two_stage_flag
+                model.model.two_stage_vanilla = two_stage_flag
+                model.roi_enable2stage = two_stage_flag
+                model.model.roi_conf_thresh = 0.1
+                with torch.no_grad():
+                    output_dicts = model.forward(**inputs.data,
+                                                 output_attentions=True,
+                                                 return_dict=False,
+                                                 output_hidden_states=False
+                                                 )
+                    base_logits = output_dicts[0][:, -1]
+
+                with torch.no_grad():
+                    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                                pad_token_id=processor.tokenizer.eos_token_id)
+
+                generated_ids = output_ids[:, input_ids_len:]
+                base_answer = \
+                    processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[
+                        0]
+
+                # --- Run with two_stage_flag = True (Two-Stage) ---
+                two_stage_flag = True
+                model.model.enable_twig = two_stage_flag
+                model.model.roi_enable2stage = two_stage_flag
+                model.model.two_stage_vanilla = two_stage_flag
+                model.roi_enable2stage = two_stage_flag
+                model.model.roi_conf_thresh = 0.1
+                with torch.no_grad():
+                    output_dicts = model.forward(**inputs.data,
+                                                 output_attentions=False,
+                                                 return_dict=False,
+                                                 output_hidden_states=False
+                                                 )
+                    two_stage_logits = output_dicts[0][:, -1]
+
+                logits_list = [base_logits, two_stage_logits]
+                KL_div_list = get_KL_divergence(logits_list)
+
+                with torch.no_grad():
+                    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, output_attentions=False,
+                                                pad_token_id=processor.tokenizer.eos_token_id)
+
+                generated_ids = output_ids[:, input_ids_len:]
+                two_stage_answer = \
+                    processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[
+                        0]
+
+                tmp_sample = {
+                    'image': image_file_str,
+                    'prompt': qs_text,  # Store the clean question
+                    'base_answer': base_answer,
+                    'two_stage_answer': two_stage_answer,
+                    'KL_divergence': KL_div_list[0],
+                    'text': ground_truth,  # This will be None for old data, or the answer
+                    'id': line.get('question_id', line.get('original_line_number', 'unknown'))
+                }
+
+                # *** MODIFIED LINE ***
+                # Use the original_line_number as the unique key
+                line_number = line['original_line_number']
+                if line_number in KL_div_analysis_dict:
+                    print(
+                        f"Worker {gpu_id}: Warning - Duplicate line_number {line_number} found. Overwriting previous entry.")
+                KL_div_analysis_dict[line_number] = tmp_sample
+
+            except Exception as e:
+                # *** MODIFIED LINE ***
+                # Update error logging to use line_number
+                print(f"Worker {gpu_id}: Error processing {line.get('original_line_number', 'unknown line')}: {e}")
+                continue
+
+        progress_bar.close()
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: A critical error occurred: {e}")
+    finally:
+        # Clean up model from GPU memory
+        del model
+        del processor
+        torch.cuda.empty_cache()
+        print(f"Worker {gpu_id}: Finished and cleaned up.")
+
+    return KL_div_analysis_dict
+
+
+
+class QwenEvalDataset(Dataset):
+    def __init__(self, questions_chunk, args):
+        self.data = questions_chunk
+        self.args = args
+        self.image_folder = args.image_folder
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        line = self.data[idx]
+        image_file_str = ''
+        ground_truth = None
+        qs_text = ""
+        image = None
+
+        try:
+            # --- Logic to parse different dataset formats ---
+            if self.args.question_file is not None and 'textvqa' in self.args.question_file or 'chartqa' in self.args.question_file:
+                qs_text = line["text"]
+                qs_text = qs_text + " Answer the question using a single word or phrase." if "Answer the question using a single word or phrase." not in qs_text else qs_text
+                image_file = line["image"]
+                image_path = os.path.join(self.image_folder, image_file)
+                image = Image.open(image_path).convert('RGB')
+                image_file_str = image_file
+                ground_truth = line.get("answers")
+
+            elif self.args.question_file is not None and 'visual_cot' in self.args.question_file:
+                qs_text = line["text"]
+                qs_text = qs_text + " Answer the question using a single word or phrase." if "Answer the question using a single word or phrase." not in qs_text else qs_text
+                image_file = line["image"]
+                image_path = os.path.join(self.image_folder, image_file)
+                image = Image.open(image_path).convert('RGB')
+                image_file_str = image_file
+                ground_truth = line.get("answer")
+
+            elif self.args.hf_dataset is not None:
+                if 'vstar' in self.args.hf_dataset:
+                    qs_text = line["text"]
+                    image = line["image"].convert('RGB')
+                    ground_truth = line.get("label")
+                    image_file_str = f"hf_dataset_{line.get('question_id', 'unknown')}"
+                elif 'POPE' in self.args.hf_dataset:
+                    qs_text = line["question"] + " Answer the question using a single word or phrase."
+                    image = line["image"].convert('RGB')
+                    ground_truth = line.get("answer")
+                    image_file_str = f"hf_dataset_{line.get('question_id', 'unknown')}"
+
+            if "<image>" in qs_text:
+                qs_text = qs_text.replace("<image>", "")  # Remove it from the text body
+
+            qs = "<image>\n" + qs_text
+
+            # Return a dict containing everything needed for inference
+            return {
+                'image': image,
+                'prompt_text': qs,
+                'meta': {
+                    'image_id': image_file_str,
+                    'prompt': qs_text,
+                    'ground_truth': ground_truth,
+                    'line_number': line['original_line_number'],
+                    'id': line.get('question_id', line.get('original_line_number', 'unknown'))
+                },
+                'valid': True
+            }
+
+        except Exception as e:
+            print(f"Error loading item {idx}: {e}")
+            # Return dummy invalid item
+            return {'valid': False, 'meta': {'line_number': line.get('original_line_number', -1)}}
+
+
+class MultiResCollator:
+    def __init__(self, model_path, resolution_bases=[64, 128, 256, 576, 768, 1024]):
+        self.model_path = model_path
+        self.resolution_bases = resolution_bases
+        self.processor = None
+
+    def __call__(self, batch):
+        if self.processor is None:
+            self.processor = AutoProcessor.from_pretrained(
+                self.model_path,
+                trust_remote_code=True,
+                padding_side='left',
+                use_fast=True
+            )
+
+        batch = [item for item in batch if item.get('valid', False)]
+        if not batch:
+            return None
+
+        # 1. Extract raw data
+        raw_images = [item['image'] for item in batch]
+        raw_prompts = [item['prompt_text'] for item in batch]
+        metas = [item['meta'] for item in batch]
+
+        inputs_by_res = {}
+
+        # 2. Process for each resolution INDEPENDENTLY
+        for res_base in self.resolution_bases:
+            # Set resolution for this pass
+            target_pixels = res_base * 28 * 28
+            self.processor.image_processor.min_pixels = target_pixels
+            self.processor.image_processor.max_pixels = target_pixels
+
+            # --- PARANOID DATA PREPARATION ---
+            # A. Deep copy images to ensure no in-place resizing affects other loops
+            current_images = [img.copy() for img in raw_images]
+
+            # B. Re-construct text from scratch for this loop.
+            # This prevents the processor from modifying the list or string states across loops.
+            current_texts = []
+            for prompt_str, img in zip(raw_prompts, current_images):
+                # Sanity check: Ensure prompt doesn't already have duplicate image tags
+                # which would confuse the processor into expecting more images.
+                clean_prompt = prompt_str
+
+                messages = [{"role": "user", "content": [{"type": "text", "text": clean_prompt}, {"image": img}]}]
+                text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+                # Qwen-VL specific replacement logic
+                text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '') \
+                    .replace('<image>', '<image><|vision_start|><|image_pad|><|vision_end|>')
+                current_texts.append(text)
+
+            try:
+                # C. Process
+                inputs = self.processor(
+                    text=current_texts,
+                    images=current_images,
+                    padding=True,
+                    return_tensors="pt"
+                )
+                inputs.data['src_images'] = current_images
+                inputs.data['processor'] = self.processor
+                inputs_by_res[res_base] = inputs
+
+            except Exception as e:
+                # Log specific error but don't crash the whole worker
+                # We save the error message to put into the results dict later
+                print(f"Error processing resolution {res_base}: {e}")
+                inputs_by_res[res_base] = None
+
+        return {
+            'metas': metas,
+            'inputs_by_res': inputs_by_res
+        }
+
+
+def check_high_res_prediction(questions_chunk, gpu_id, args):
+    device = f'cuda:{gpu_id}'
+    ms_answer_analysis_dict = {}
+    max_new_tokens = 128
+    resolution_bases = [64, 128, 256, 576, 768, 1024]
+    from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
+
+    # --- SETUP DATALOADER ---
+    batch_size = getattr(args, 'batch_size', 1)
+    num_workers = getattr(args, 'num_workers', 4)
+    image_root_folder = getattr(args, 'image_folder', '/home/yuheng/datasets')
+    dataset = QwenEvalDataset(questions_chunk, args)
+    collator = MultiResCollator(args.model_path, resolution_bases)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collator,
+        pin_memory=True  # Faster transfer to GPU
+    )
+
+    try:
+        # Load Model (No processor needed here anymore!)
+        print(f"Worker {gpu_id}: Loading model...")
+        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device
+        )
+        model.eval()
+
+        # Flags
+        roi_flag = False
+        high_res_flag = True
+        model.model.enable_twig = high_res_flag or roi_flag
+        model.model.roi_enable2stage = roi_flag
+        model.model.two_stage_vanilla = roi_flag
+        model.roi_enable2stage = roi_flag
+        model.model.is_debug = True
+
+        tokenizer = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True).tokenizer
+
+        print(f"Worker {gpu_id}: Starting inference...")
+
+        # --- MAIN INFERENCE LOOP ---
+        progress_bar = tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id)
+
+        for batch_data in progress_bar:
+            metas = batch_data['metas']
+            inputs_by_res = batch_data['inputs_by_res']
+            # Temporary storage for this batch
+            batch_results = {i: {} for i in range(len(metas))}
+            batch_pred = {i: {} for i in range(len(metas))}
+            # Loop through pre-processed resolutions
+            for res_base in resolution_bases:
+                inputs = inputs_by_res.get(res_base)
+                if inputs is None:
+                    # Handle preprocessing failure
+                    for idx in range(len(metas)): batch_results[idx][res_base] = "ERROR_PREPROC"
+                    continue
+
+                # Move pre-processed tensors to GPU
+                inputs = inputs.to(device)
+                with torch.no_grad():
+                    output_ids = model.generate(
+                        **inputs,
+                        max_new_tokens=max_new_tokens,
+                        output_attentions=False,
+                        pad_token_id=tokenizer.eos_token_id
+                    )
+
+                # Decode
+                input_ids_len = inputs.input_ids.shape[1]
+                generated_ids = output_ids[:, input_ids_len:]
+                batch_answers = tokenizer.batch_decode(
+                    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+                )
+
+                for idx, answer in enumerate(batch_answers):
+                    batch_results[idx][res_base] = answer
+                    batch_pred[idx][res_base] = model.model.high_res_pred.flatten().float().cpu()
+                # Cleanup GPU memory for this resolution immediately
+                del inputs, output_ids, generated_ids
+
+            # --- Show RESULTS ---
+            for idx, meta in enumerate(metas):
+                line_number = meta['line_number']
+                ordered_answers = [batch_results[idx].get(r, "ERROR") for r in resolution_bases]
+
+                tmp_sample = {
+                    'image': meta['image_id'],
+                    'prompt': meta['prompt'],
+                    'text': meta['ground_truth'],
+                    'id': meta['id'],
+                    'res': resolution_bases,
+                    'ms_answers': ordered_answers
+                }
+                ms_answer_analysis_dict[line_number] = tmp_sample
+
+            plot_single_sample(ms_answer_analysis_dict[line_number], image_root_folder)
+            pass
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: Critical Error: {e}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        torch.cuda.empty_cache()
+
+def worker_process_ms(questions_chunk, gpu_id, args):
+    device = f'cuda:{gpu_id}'
+    ms_answer_analysis_dict = {}
+    max_new_tokens = 128
+    resolution_bases = [64, 128, 256, 576, 768, 1024]
+
+    # --- SETUP DATALOADER ---
+    batch_size = getattr(args, 'batch_size', 8)
+    # Use 2-4 workers to ensure CPU preprocessing keeps up with GPU
+    num_workers = getattr(args, 'num_workers', 4)
+
+    dataset = QwenEvalDataset(questions_chunk, args)
+
+    # Initialize our custom collator
+    collator = MultiResCollator(args.model_path, resolution_bases)
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collator,
+        pin_memory=True  # Faster transfer to GPU
+    )
+
+    if len(dataset) == 0: return ms_answer_analysis_dict
+
+    try:
+        # Load Model (No processor needed here anymore!)
+        print(f"Worker {gpu_id}: Loading model...")
+        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device
+        )
+        model.eval()
+
+        # Flags
+        two_stage_flag = False
+        model.model.enable_twig = two_stage_flag
+        model.model.roi_enable2stage = two_stage_flag
+        model.model.two_stage_vanilla = two_stage_flag
+        model.roi_enable2stage = two_stage_flag
+
+        tokenizer = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True).tokenizer
+
+        print(f"Worker {gpu_id}: Starting inference...")
+
+        # --- MAIN INFERENCE LOOP ---
+        progress_bar = tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id)
+
+        for batch_data in progress_bar:
+            if batch_data is None: continue
+
+            metas = batch_data['metas']
+            inputs_by_res = batch_data['inputs_by_res']
+
+            # Temporary storage for this batch
+            batch_results = {i: {} for i in range(len(metas))}
+
+            # Loop through pre-processed resolutions
+            for res_base in resolution_bases:
+                inputs = inputs_by_res.get(res_base)
+
+                if inputs is None:
+                    # Handle preprocessing failure
+                    for idx in range(len(metas)): batch_results[idx][res_base] = "ERROR_PREPROC"
+                    continue
+
+                try:
+                    # Move pre-processed tensors to GPU
+                    inputs = inputs.to(device)
+
+                    with torch.no_grad():
+                        output_ids = model.generate(
+                            **inputs,
+                            max_new_tokens=max_new_tokens,
+                            output_attentions=False,
+                            pad_token_id=tokenizer.eos_token_id
+                        )
+
+                    # Decode
+                    input_ids_len = inputs.input_ids.shape[1]
+                    generated_ids = output_ids[:, input_ids_len:]
+                    batch_answers = tokenizer.batch_decode(
+                        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+                    )
+
+                    for idx, answer in enumerate(batch_answers):
+                        batch_results[idx][res_base] = answer
+
+                except RuntimeError as e:
+                    msg = "ERROR_OOM" if "out of memory" in str(e) else "ERROR_RUNTIME"
+                    print(f"Worker {gpu_id}: {msg} at res {res_base}")
+                    torch.cuda.empty_cache()
+                    for idx in range(len(metas)): batch_results[idx][res_base] = msg
+
+                # Cleanup GPU memory for this resolution immediately
+                del inputs, output_ids, generated_ids
+
+            # --- SAVE RESULTS ---
+            for idx, meta in enumerate(metas):
+                line_number = meta['line_number']
+                ordered_answers = [batch_results[idx].get(r, "ERROR") for r in resolution_bases]
+
+                tmp_sample = {
+                    'image': meta['image_id'],
+                    'prompt': meta['prompt'],
+                    'text': meta['ground_truth'],
+                    'id': meta['id'],
+                    'res': resolution_bases,
+                    'ms_answers': ordered_answers
+                }
+                ms_answer_analysis_dict[line_number] = tmp_sample
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: Critical Error: {e}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        torch.cuda.empty_cache()
+
+    return ms_answer_analysis_dict
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Parallel KL Divergence Evaluation")
+    parser.add_argument('--model-path', type=str, required=True,
+                        help="Path to the model and processor")
+    parser.add_argument('--question-file', type=str, default=None,
+                        help="Path to the JSONL question file (E.g., TextVQA)")
+    parser.add_argument('--image-folder', type=str, default=None,
+                        help="Path to the image directory (for JSONL file)")
+    parser.add_argument('--dst-save-path', type=str, required=True,
+                        help="Path to save the output .pkl file")
+    parser.add_argument('--gpu-ids', type=str, default='0,1',
+                        help="Comma-separated list of GPU IDs to use (e.g., '0,1,2,3')")
+    # --- ADDED: New HF Dataset args ---
+    parser.add_argument('--hf_dataset', type=str, default=None,
+                        help="Name of the Hugging Face dataset (e.g., 'craigwu/vstar_bench')")
+    parser.add_argument('--dataset_split', type=str, default='test',
+                        help="Split to use from the Hugging Face dataset (e.g., 'test')")
+    parser.add_argument('--is_check_pred', type=bool, default=False,
+                        help="Whether to check high-res predictions")
+    ### --- END ADDITION ---
+    args = parser.parse_args()
+
+
+    # --- 1. Define a Wrapper Function ---
+    # This helper must be defined at the top level (global scope) so it can be pickled.
+    # --- 2. The Main Execution Block ---
+    # Ensure start method is spawn (Crucial for PyTorch+CUDA)
+    try:
+        mp.set_start_method('spawn', force=True)
+    except RuntimeError:
+        pass
+
+    # 1. Prepare Data and Split into Chunks
+    gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
+    num_gpus = len(gpu_ids)
+    questions = []
+
+    if args.hf_dataset:
+        print(f"Loading Hugging Face dataset {args.hf_dataset}, split {args.dataset_split}")
+        from datasets import load_dataset
+
+        dataset = load_dataset(args.hf_dataset, split=args.dataset_split)
+        print("Casting 'image' column to load from cache...")
+        questions_list = list(dataset)
+
+        new_questions = []
+        for i, tmp_ques in enumerate(questions_list):
+            tmp_ques['original_line_number'] = i
+            new_questions.append(tmp_ques)
+        questions = new_questions
+        print(f"Loaded {len(questions)} preprocessed questions from Hugging Face.")
+
+    elif args.question_file:
+        print("Loading and preprocessing questions from JSONL file...")
+        if not args.image_folder:
+            raise ValueError("--image-folder is required when using --question-file")
+
+        questions_list = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+
+        new_questions = []
+        for i, tmp_ques in enumerate(questions_list):
+            qs = tmp_ques['text'] if 'text' in tmp_ques else tmp_ques['prompt']
+            if 'Reference OCR' in qs:
+                qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
+            tmp_ques['text'] = qs
+            tmp_ques['original_line_number'] = i
+            new_questions.append(tmp_ques)
+        questions = new_questions
+        print(f"Loaded {len(questions)} preprocessed questions from local file.")
+
+    else:
+        raise ValueError("You must provide either --question-file or --hf_dataset")
+
+    # Create chunks
+    chunk_size = math.ceil(len(questions) / num_gpus)
+    chunks = [questions[i:i + chunk_size] for i in range(0, len(questions), chunk_size)]
+
+    # 2. Run processes in parallel using mp.Process (Manual Management)
+    print(f"Starting parallel processing on {num_gpus} GPUs: {gpu_ids}")
+
+    # Create a Manager to hold the results safely across processes
+    manager = mp.Manager()
+    return_dict = manager.dict()
+    processes = []
+
+    for i in range(num_gpus):
+        chunk_data = chunks[i] if i < len(chunks) else []
+
+        # We pass the shared 'return_dict' to the wrapper
+        p = mp.Process(
+            target=worker_wrapper,
+            args=(chunk_data, gpu_ids[i], args, return_dict)
+        )
+        p.start()
+        processes.append(p)
+
+    # Wait for all processes to finish
+    for p in processes:
+        p.join()
+
+    # 3. Gather and Save Results
+    print("All chunks processed. Gathering results...")
+    final_ms_dict = {}
+
+    # Extract results from the shared return_dict
+    for gpu_id, sub_dict in return_dict.items():
+        final_ms_dict.update(sub_dict)
+
+    with open(os.path.expanduser(args.dst_save_path), 'wb') as f:
+        pickle.dump(final_ms_dict, f)
+
+    print(f"✅ Successfully processed and saved {len(final_ms_dict)} multi-resolution answers to {args.dst_save_path}.")
+
+
+'''
+CUDA_VISIBLE_DEVICES=2,3 python analysis/making_gated_training_samples.py \
+    --model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+    --question-file /home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl \
+    --image-folder ../llava/playground/data/eval/textvqa/train_images \
+    --dst-save-path /home/yuheng/datasets/gated_roi_datasets/roi_KLdiv_textvqa_576res.pkl \
+    --gpu-ids 0,1
+
+CUDA_VISIBLE_DEVICES=3 python analysis/making_gated_training_samples.py \
+    --model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+    --question-file /home/yuheng/datasets/visual_cot_llava_subset68k.jsonl \
+    --image-folder /home/yuheng/datasets \
+    --dst-save-path /home/yuheng/datasets/gated_roi_datasets/roi_KLdiv_vcot68k_576res.pkl \
+    --gpu-ids 0
+
+python analysis/making_gated_training_samples.py \
+    --model-path /root/public/Qwen/Qwen2.5-VL-7B-Instruct \
+    --question-file /root/epfs/sproject/datasets/visual_cot_docvqa_subset33k.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/gated_roi_datasets/roi_vcot_docvqa33k_ms.pkl \
+    --gpu-ids 0,1
+
+
+python analysis/making_gated_training_samples.py \
+    --model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+    --question-file /home/yuheng/datasets/visual_cot_llava_subset68k.jsonl \
+    --image-folder /home/yuheng/datasets \
+    --dst-save-path /home/yuheng/datasets/gated_roi_datasets/roi_KLdiv_vcot68k_576res_wprediction.pkl \
+    --gpu-ids 0,1,2,3
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 python analysis/making_gated_training_samples.py \
+--model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+--hf_dataset "lmms-lab/vstar-bench" --dataset_split "test" \
+--dst-save-path ./excluded/vstar_analysis.pkl --gpu-ids 0,1,2,3
+
+CUDA_VISIBLE_DEVICES=0,1,2,3 python analysis/making_gated_training_samples.py \
+--model-path /home/yuheng/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K16T3-152k-v1bf16Mheads-twiginit-filled \
+--hf_dataset "lmms-lab/POPE" --dataset_split "test" \
+--dst-save-path ./excluded/pope_analysis.pkl --gpu-ids 0,1,2,3
+
+python analysis/making_gated_training_samples.py \
+    --model-path /root/public/Qwen/Qwen2.5-VL-7B-Instruct \
+    --question-file /root/epfs/sproject/datasets/chartqa_28k_qa.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/roi_chartqa28k_ms.pkl \
+    --gpu-ids 0,1,2,3
+
+'''

QWENVL-PRIVATE/analysis/making_roi_model_response.py

@@ -0,0 +1,565 @@
+import sys
+import os
+script_dir = os.path.dirname(os.path.abspath(__file__))
+project_root = os.path.dirname(script_dir)
+sys.path.append(project_root)
+
+from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
+#from qwen_src.qwen2_vl import Qwen2VLForConditionalGeneration
+from qwen_src.qwen3_vl.modeling_qwen3_vl_batch import Qwen3VLForConditionalGeneration
+import torch
+import matplotlib.pyplot as plt
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor
+import math
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+
+import json
+import os
+from tqdm import tqdm
+
+import pickle
+import argparse
+import multiprocessing as mp
+import pandas as pd  # <-- Make sure to import pandas
+from matching_utils import TextVQAAccuracyEvaluator
+from check_gate_label import plot_single_sample
+
+
+def infer_model_type(model_path, model_type):
+    if model_type and model_type != "auto":
+        return model_type
+    lower = model_path.lower()
+    if "qwen3" in lower:
+        return "qwen3"
+    if "qwen2.5" in lower or "qwen2_5" in lower or "qwen2-5" in lower:
+        return "qwen2_5"
+    if "qwen2" in lower:
+        return "qwen2"
+    return "qwen2_5"
+
+
+def resolve_processor_path(model_path, model_type, processor_path=None):
+    if processor_path:
+        return processor_path
+    lower = model_path.lower()
+    public_prefix = "/root/public/Qwen"
+
+    def to_repo_path(name):
+        if "/root/public" in model_path:
+            return f"{public_prefix}/{name}"
+        return f"Qwen/{name}"
+
+    if model_type == "qwen3":
+        if "8b" in lower:
+            return to_repo_path("Qwen3-VL-8B-Instruct")
+        if "4b" in lower:
+            return to_repo_path("Qwen3-VL-4B-Instruct")
+    if model_type == "qwen2_5":
+        if "7b" in lower:
+            return to_repo_path("Qwen2.5-VL-7B-Instruct")
+        if "3b" in lower:
+            return to_repo_path("Qwen2.5-VL-3B-Instruct")
+    if model_type == "qwen2" and "7b" in lower:
+        return to_repo_path("Qwen2-VL-7B-Instruct")
+
+    return model_path
+
+
+def resolve_model_class(model_type):
+    if model_type == "qwen3":
+        if Qwen3VLForConditionalGeneration is None:
+            raise RuntimeError("Qwen3VLForConditionalGeneration not available.")
+        return Qwen3VLForConditionalGeneration
+    if model_type == "qwen2":
+        return Qwen2VLForConditionalGeneration
+    return Qwen2_5_VLForConditionalGeneration
+
+
+def set_twig_flags(model, model_type, two_stage_flag):
+    if model_type == "qwen3":
+        if hasattr(model, "enable_twig"):
+            model.enable_twig = two_stage_flag
+        if hasattr(model, "roi_enable2stage"):
+            model.roi_enable2stage = two_stage_flag
+        if hasattr(model, "language_model"):
+            if hasattr(model.language_model, "enable_twig"):
+                model.language_model.enable_twig = two_stage_flag
+            if hasattr(model.language_model, "roi_enable2stage"):
+                model.language_model.roi_enable2stage = two_stage_flag
+            if hasattr(model.language_model, "roi_conf_thresh"):
+                model.language_model.roi_conf_thresh = 0.15
+        return
+
+    if hasattr(model, "model"):
+        if hasattr(model.model, "enable_twig"):
+            model.model.enable_twig = two_stage_flag
+        if hasattr(model.model, "roi_enable2stage"):
+            model.model.roi_enable2stage = two_stage_flag
+        if hasattr(model.model, "two_stage_vanilla"):
+            model.model.two_stage_vanilla = two_stage_flag
+        if hasattr(model.model, "roi_conf_thresh"):
+            model.model.roi_conf_thresh = 0.15
+    if hasattr(model, "roi_enable2stage"):
+        model.roi_enable2stage = two_stage_flag
+def get_KL_divergence(logits_list,):
+    import torch.nn as nn
+    from torch.nn import functional as F
+    KL_div = nn.KLDivLoss(reduction='batchmean')
+    KL_div_list = [KL_div(F.log_softmax(logits_list[-1], dim=1), F.softmax(logits_list[i], dim=1)).item() for i in range(len(logits_list))]
+    return KL_div_list
+
+
+from torch.utils.data import Dataset, DataLoader
+
+
+def worker_wrapper(chunk, gpu_id, args, return_dict):
+    """
+    Wraps the main worker function to handle return values via shared dictionary.
+    """
+
+    result = worker_process(chunk, gpu_id, args)
+    # Store result in the shared dictionary using gpu_id as key
+    return_dict[gpu_id] = result
+
+
+
+class QwenEvalDataset(Dataset):
+    def __init__(self, questions_chunk, args):
+        self.data = questions_chunk
+        self.args = args
+        self.image_folder = args.image_folder
+        self.use_short_prompt = getattr(args, 'use_short_prompt', False)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        line = self.data[idx]
+        image_file_str = ''
+        ground_truth = None
+        qs_text = ""
+        image = None
+
+        try:
+
+            qs_text = line["text"]
+            qs_text = qs_text
+            image_file = line["image"]
+            image_path = os.path.join(self.image_folder, image_file)
+            image = Image.open(image_path).convert('RGB')
+            image_file_str = image_file
+            ground_truth = line.get("answer") if "answer" in line else line.get("answers", None)
+
+            if "<image>" in qs_text:
+                qs_text = qs_text.replace("<image>", "")  # Remove it from the text body
+
+            qs = "<image>\n" + qs_text
+
+            if self.use_short_prompt:
+                qs = qs + "\nAnswer the question using a single word or phrase."
+
+            # Return a dict containing everything needed for inference
+            return {
+                'image': image,
+                'prompt_text': qs,
+                'meta': {
+                    'image_id': image_file_str,
+                    'prompt': qs_text,
+                    'ground_truth': ground_truth,
+                    'line_number': line['original_line_number'],
+                    'id': line.get('question_id', line.get('original_line_number', 'unknown'))
+                },
+                'valid': True
+            }
+
+        except Exception as e:
+            print(f"Error loading item {idx}: {e}")
+            # Return dummy invalid item
+            return {'valid': False, 'meta': {'line_number': line.get('original_line_number', -1)}}
+
+
+class SingleResCollator:
+    def __init__(self, processor_path, patch_size):
+        self.processor_path = processor_path
+        self.processor = None
+        # Define the requested limits
+        self.min_pixels = 128 * patch_size * patch_size
+        self.max_pixels = 576 * patch_size * patch_size
+
+    def __call__(self, batch):
+        if self.processor is None:
+            # Initialize processor once per worker
+            self.processor = AutoProcessor.from_pretrained(
+                self.processor_path,
+                trust_remote_code=True,
+                padding_side='left',
+                use_fast=True
+            )
+            # Set the resolution constraints globally for this processor
+            self.processor.image_processor.min_pixels = self.min_pixels
+            self.processor.image_processor.max_pixels = self.max_pixels
+
+
+        batch = [item for item in batch if item.get('valid', False)]
+        if not batch:
+            return None
+
+        # 1. Extract raw data
+        raw_images = [item['image'] for item in batch]
+        raw_prompts = [item['prompt_text'] for item in batch]
+        metas = [item['meta'] for item in batch]
+
+        # 2. Prepare text and images
+        current_texts = []
+        for prompt_str, img in zip(raw_prompts, raw_images):
+            # Qwen-VL-Chat template construction
+            messages = [{"role": "user", "content": [{"type": "text", "text": prompt_str}, {"image": img}]}]
+            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+            # Specific replacement logic for Qwen2.5-VL
+            text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '') \
+                .replace('<image>', '<image><|vision_start|><|image_pad|><|vision_end|>')
+            current_texts.append(text)
+
+        try:
+            # 3. Process Batch (Single Resolution Setting)
+            inputs = self.processor(
+                text=current_texts,
+                images=raw_images,
+                padding=True,
+                return_tensors="pt"
+            )
+            # Attach extra data needed later if necessary (or just return inputs)
+            inputs.data['src_images'] = raw_images
+
+            return {
+                'metas': metas,
+                'inputs': inputs
+            }
+
+        except Exception as e:
+            print(f"Error processing batch: {e}")
+            return None
+
+
+def worker_process(questions_chunk, gpu_id, args):
+    device = f'cuda:{gpu_id}'
+    answer_analysis_dict = {}  # Renamed from ms_answer_analysis_dict
+    max_new_tokens = 128
+
+    # --- SETUP DATALOADER ---
+    batch_size = getattr(args, 'batch_size', 1)
+    num_workers = getattr(args, 'num_workers', 4)
+
+    dataset = QwenEvalDataset(questions_chunk, args)
+
+    # Initialize the NEW single resolution collator
+    collator = SingleResCollator(args.processor_path, args.patch_size)
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collator,
+        pin_memory=True
+    )
+
+    if len(dataset) == 0: return answer_analysis_dict
+
+    try:
+        # Load Model
+        print(f"Worker {gpu_id}: Loading model...")
+        model_class = resolve_model_class(args.model_type)
+        model = model_class.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device
+        )
+        model.eval()
+
+        # Flags (keeping your original config)
+        two_stage_flag = True
+        set_twig_flags(model, args.model_type, two_stage_flag)
+
+        # Tokenizer for decoding
+        tokenizer = AutoProcessor.from_pretrained(args.processor_path, trust_remote_code=True).tokenizer
+
+        print(f"Worker {gpu_id}: Starting inference...")
+
+        # --- MAIN INFERENCE LOOP ---
+        progress_bar = tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id)
+
+        for batch_data in progress_bar:
+            if batch_data is None: continue
+
+            metas = batch_data['metas']
+            inputs = batch_data['inputs']  # Now getting direct inputs, not inputs_by_res
+
+            try:
+                # Move to GPU
+                inputs = inputs.to(device)
+                batch_size = inputs.input_ids.shape[0]
+                assert batch_size == 1, "Batch size greater than 1 is not supported in this script."
+                inputs.data['processor'] = collator.processor
+                with torch.no_grad():
+                    output_ids = model.generate(
+                        **inputs,
+                        max_new_tokens=max_new_tokens,
+                        output_attentions=False,
+                        pad_token_id=tokenizer.eos_token_id
+                    )
+
+                # Decode
+                input_ids_len = inputs.input_ids.shape[1]
+                generated_ids = output_ids[:, input_ids_len:]
+                batch_answers = tokenizer.batch_decode(
+                    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+                )
+                surrounding_bbox = model.model.surrounding_bbox if hasattr(model.model, 'surrounding_bbox') else None
+                roi_mask = model.model.roi_mask if hasattr(model.model, 'roi_mask') else None
+                if surrounding_bbox is not None and roi_mask is not None:
+                    ###convert them to cpu numpy
+                    surrounding_bbox = surrounding_bbox.cpu().numpy()
+                    roi_mask = roi_mask.cpu().numpy()
+
+                # --- SAVE RESULTS ---
+                for idx, (meta, answer) in enumerate(zip(metas, batch_answers)):
+                    line_number = meta['line_number']
+
+                    tmp_sample = {
+                        'image': meta['image_id'],
+                        'prompt': meta['prompt'],
+                        'text': meta['ground_truth'],
+                        'id': meta['id'],
+                        'response': answer,  # Storing single response
+                        'surrounding_bbox': surrounding_bbox,
+                        'roi_mask': roi_mask
+                    }
+                    answer_analysis_dict[line_number] = tmp_sample
+
+                # Cleanup GPU memory for this batch
+                del inputs, output_ids, generated_ids
+
+            except RuntimeError as e:
+                msg = "ERROR_OOM" if "out of memory" in str(e) else f"ERROR_RUNTIME: {e}"
+                print(f"Worker {gpu_id}: {msg}")
+                torch.cuda.empty_cache()
+                # Log errors for this batch
+                for meta in metas:
+                    line_number = meta['line_number']
+                    answer_analysis_dict[line_number] = {
+                        'id': meta['id'],
+                        'response': msg
+                    }
+
+    except Exception as e:
+        print(f"Worker {gpu_id}: Critical Error: {e}")
+        import traceback
+        traceback.print_exc()
+    finally:
+        torch.cuda.empty_cache()
+
+    return answer_analysis_dict
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Parallel KL Divergence Evaluation")
+    parser.add_argument('--model-path', type=str, required=True,
+                        help="Path to the model and processor")
+    parser.add_argument('--model-type', type=str, default="auto",
+                        choices=["auto", "qwen2_5", "qwen2", "qwen3"],
+                        help="Model family for loading (auto detects from model path).")
+    parser.add_argument('--processor-path', type=str, default="",
+                        help="Override processor path if auto detection is incorrect.")
+    parser.add_argument('--question-file', type=str, default=None,
+                        help="Path to the JSONL question file (E.g., TextVQA)")
+    parser.add_argument('--image-folder', type=str, default=None,
+                        help="Path to the image directory (for JSONL file)")
+    parser.add_argument('--dst-save-path', type=str, required=True,
+                        help="Path to save the output .pkl file")
+    parser.add_argument('--gpu-ids', type=str, default='0,1',
+                        help="Comma-separated list of GPU IDs to use (e.g., '0,1,2,3')")
+    # --- ADDED: New HF Dataset args ---
+    parser.add_argument('--hf_dataset', type=str, default=None,
+                        help="Name of the Hugging Face dataset (e.g., 'craigwu/vstar_bench')")
+    parser.add_argument('--dataset_split', type=str, default='test',
+                        help="Split to use from the Hugging Face dataset (e.g., 'test')")
+    parser.add_argument('--is_check_pred', type=bool, default=False,
+                        help="Whether to check high-res predictions")
+    ### --- END ADDITION ---
+    parser.add_argument('--use-short-prompt', action='store_true',
+                        help="Whether to use short prompt format")
+    args = parser.parse_args()
+    args.model_type = infer_model_type(args.model_path, args.model_type)
+    args.processor_path = resolve_processor_path(
+        args.model_path, args.model_type, args.processor_path
+    )
+    args.patch_size = 32 if args.model_type == "qwen3" else 28
+
+
+    # --- 1. Define a Wrapper Function ---
+    # This helper must be defined at the top level (global scope) so it can be pickled.
+    # --- 2. The Main Execution Block ---
+    # Ensure start method is spawn (Crucial for PyTorch+CUDA)
+    try:
+        mp.set_start_method('spawn', force=True)
+    except RuntimeError:
+        pass
+
+    # 1. Prepare Data and Split into Chunks
+    gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
+    num_gpus = len(gpu_ids)
+    questions = []
+
+    if args.hf_dataset:
+        print(f"Loading Hugging Face dataset {args.hf_dataset}, split {args.dataset_split}")
+        from datasets import load_dataset
+
+        dataset = load_dataset(args.hf_dataset, split=args.dataset_split)
+        print("Casting 'image' column to load from cache...")
+        questions_list = list(dataset)
+
+        new_questions = []
+        for i, tmp_ques in enumerate(questions_list):
+            tmp_ques['original_line_number'] = i
+            new_questions.append(tmp_ques)
+        questions = new_questions
+        print(f"Loaded {len(questions)} preprocessed questions from Hugging Face.")
+
+    elif args.question_file:
+        print("Loading and preprocessing questions from JSONL file...")
+        if not args.image_folder:
+            raise ValueError("--image-folder is required when using --question-file")
+
+        questions_list = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+
+        new_questions = []
+        for i, tmp_ques in enumerate(questions_list):
+            qs = tmp_ques['text'] if 'text' in tmp_ques else tmp_ques['prompt']
+            if 'Reference OCR' in qs:
+                qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
+            tmp_ques['text'] = qs
+            tmp_ques['original_line_number'] = i
+            new_questions.append(tmp_ques)
+        questions = new_questions
+        print(f"Loaded {len(questions)} preprocessed questions from local file.")
+
+    else:
+        raise ValueError("You must provide either --question-file or --hf_dataset")
+
+    # Create chunks
+    chunk_size = math.ceil(len(questions) / num_gpus)
+    chunks = [questions[i:i + chunk_size] for i in range(0, len(questions), chunk_size)]
+
+    # 2. Run processes in parallel using mp.Process (Manual Management)
+    print(f"Starting parallel processing on {num_gpus} GPUs: {gpu_ids}")
+
+    # Create a Manager to hold the results safely across processes
+    manager = mp.Manager()
+    return_dict = manager.dict()
+    processes = []
+
+    for i in range(num_gpus):
+        chunk_data = chunks[i] if i < len(chunks) else []
+
+        # We pass the shared 'return_dict' to the wrapper
+        p = mp.Process(
+            target=worker_wrapper,
+            args=(chunk_data, gpu_ids[i], args, return_dict)
+        )
+        p.start()
+        processes.append(p)
+
+    # Wait for all processes to finish
+    for p in processes:
+        p.join()
+
+    # 3. Gather and Save Results
+    print("All chunks processed. Gathering results...")
+    final_dict = {}
+
+    # Extract results from the shared return_dict
+    for gpu_id, sub_dict in return_dict.items():
+        final_dict.update(sub_dict)
+
+    #check save directory
+    os.makedirs(os.path.dirname(os.path.expanduser(args.dst_save_path)), exist_ok=True)
+    with open(os.path.expanduser(args.dst_save_path), 'wb') as f:
+        pickle.dump(final_dict, f)
+
+    print(f"✅ Successfully processed and saved {len(final_dict)} answers to {args.dst_save_path}.")
+
+
+'''
+CUDA_VISIBLE_DEVICES=2,3 python analysis/making_roi_model_response.py \
+    --model-path ./output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl \
+    --image-folder ../llava/playground/data/eval/textvqa/train_images \
+    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B-roi.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+
+CUDA_VISIBLE_DEVICES=2,3 python analysis/making_roi_model_response.py \
+    --model-path ./output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /home/yuheng/datasets/vstar_coco_spatial_relation_data.jsonl \
+    --image-folder /home/yuheng/datasets \
+    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B-roi.pkl \
+    --gpu-ids 0,1
+
+
+
+python analysis/making_roi_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/textvqa/converted_llava_style_train.jsonl \
+    --image-folder /root/epfs/sproject/datasets/textvqa/train_images \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B-roi-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+    
+python analysis/making_roi_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/chartqa_28k_qa.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/chartqa_qwen2.5vl-7B-roi-short-prompt.pkl \\
+    --use-short-prompt \
+    --gpu-ids 0,1
+
+python analysis/making_roi_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/visual_cot_infovqa_subset15k.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_infovqa_qwen2.5vl-7B-roi-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+    
+python analysis/making_roi_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/visual_cot_docvqa_subset33k.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_docvqa_qwen2.5vl-7B-roi-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+    
+python analysis/making_roi_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/vstar_coco_spatial_relation_data.jsonl \
+    --image-folder /root/epfs/sproject/datasets/ \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B-roi-short-prompt.pkl \
+    --use-short-prompt \
+    --gpu-ids 0,1
+
+
+
+
+python analysis/making_base_model_response.py \
+    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
+    --question-file /root/epfs/sproject/datasets/vstar_coco_spatial_relation_data.jsonl \
+    --image-folder /root/epfs/sproject/datasets \
+    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B-roi-short-prompt.pkl \
+    --gpu-ids 0,1,2,3
+'''

QWENVL-PRIVATE/analysis/matching_utils.py

@@ -0,0 +1,257 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import re
+
+from tqdm import tqdm
+
+
+class EvalAIAnswerProcessor:
+    """
+    Processes an answer similar to Eval AI
+        copied from
+        https://github.com/facebookresearch/mmf/blob/c46b3b3391275b4181567db80943473a89ab98ab/pythia/tasks/processors.py#L897
+    """
+
+    CONTRACTIONS = {
+        "aint": "ain't",
+        "arent": "aren't",
+        "cant": "can't",
+        "couldve": "could've",
+        "couldnt": "couldn't",
+        "couldn'tve": "couldn't've",
+        "couldnt've": "couldn't've",
+        "didnt": "didn't",
+        "doesnt": "doesn't",
+        "dont": "don't",
+        "hadnt": "hadn't",
+        "hadnt've": "hadn't've",
+        "hadn'tve": "hadn't've",
+        "hasnt": "hasn't",
+        "havent": "haven't",
+        "hed": "he'd",
+        "hed've": "he'd've",
+        "he'dve": "he'd've",
+        "hes": "he's",
+        "howd": "how'd",
+        "howll": "how'll",
+        "hows": "how's",
+        "Id've": "I'd've",
+        "I'dve": "I'd've",
+        "Im": "I'm",
+        "Ive": "I've",
+        "isnt": "isn't",
+        "itd": "it'd",
+        "itd've": "it'd've",
+        "it'dve": "it'd've",
+        "itll": "it'll",
+        "let's": "let's",
+        "maam": "ma'am",
+        "mightnt": "mightn't",
+        "mightnt've": "mightn't've",
+        "mightn'tve": "mightn't've",
+        "mightve": "might've",
+        "mustnt": "mustn't",
+        "mustve": "must've",
+        "neednt": "needn't",
+        "notve": "not've",
+        "oclock": "o'clock",
+        "oughtnt": "oughtn't",
+        "ow's'at": "'ow's'at",
+        "'ows'at": "'ow's'at",
+        "'ow'sat": "'ow's'at",
+        "shant": "shan't",
+        "shed've": "she'd've",
+        "she'dve": "she'd've",
+        "she's": "she's",
+        "shouldve": "should've",
+        "shouldnt": "shouldn't",
+        "shouldnt've": "shouldn't've",
+        "shouldn'tve": "shouldn't've",
+        "somebody'd": "somebodyd",
+        "somebodyd've": "somebody'd've",
+        "somebody'dve": "somebody'd've",
+        "somebodyll": "somebody'll",
+        "somebodys": "somebody's",
+        "someoned": "someone'd",
+        "someoned've": "someone'd've",
+        "someone'dve": "someone'd've",
+        "someonell": "someone'll",
+        "someones": "someone's",
+        "somethingd": "something'd",
+        "somethingd've": "something'd've",
+        "something'dve": "something'd've",
+        "somethingll": "something'll",
+        "thats": "that's",
+        "thered": "there'd",
+        "thered've": "there'd've",
+        "there'dve": "there'd've",
+        "therere": "there're",
+        "theres": "there's",
+        "theyd": "they'd",
+        "theyd've": "they'd've",
+        "they'dve": "they'd've",
+        "theyll": "they'll",
+        "theyre": "they're",
+        "theyve": "they've",
+        "twas": "'twas",
+        "wasnt": "wasn't",
+        "wed've": "we'd've",
+        "we'dve": "we'd've",
+        "weve": "we've",
+        "werent": "weren't",
+        "whatll": "what'll",
+        "whatre": "what're",
+        "whats": "what's",
+        "whatve": "what've",
+        "whens": "when's",
+        "whered": "where'd",
+        "wheres": "where's",
+        "whereve": "where've",
+        "whod": "who'd",
+        "whod've": "who'd've",
+        "who'dve": "who'd've",
+        "wholl": "who'll",
+        "whos": "who's",
+        "whove": "who've",
+        "whyll": "why'll",
+        "whyre": "why're",
+        "whys": "why's",
+        "wont": "won't",
+        "wouldve": "would've",
+        "wouldnt": "wouldn't",
+        "wouldnt've": "wouldn't've",
+        "wouldn'tve": "wouldn't've",
+        "yall": "y'all",
+        "yall'll": "y'all'll",
+        "y'allll": "y'all'll",
+        "yall'd've": "y'all'd've",
+        "y'alld've": "y'all'd've",
+        "y'all'dve": "y'all'd've",
+        "youd": "you'd",
+        "youd've": "you'd've",
+        "you'dve": "you'd've",
+        "youll": "you'll",
+        "youre": "you're",
+        "youve": "you've",
+    }
+
+    NUMBER_MAP = {
+        "none": "0",
+        "zero": "0",
+        "one": "1",
+        "two": "2",
+        "three": "3",
+        "four": "4",
+        "five": "5",
+        "six": "6",
+        "seven": "7",
+        "eight": "8",
+        "nine": "9",
+        "ten": "10",
+    }
+    ARTICLES = ["a", "an", "the"]
+    PERIOD_STRIP = re.compile(r"(?!<=\d)(\.)(?!\d)")
+    COMMA_STRIP = re.compile(r"(?<=\d)(\,)+(?=\d)")
+    PUNCTUATIONS = [
+        ";",
+        r"/",
+        "[",
+        "]",
+        '"',
+        "{",
+        "}",
+        "(",
+        ")",
+        "=",
+        "+",
+        "\\",
+        "_",
+        "-",
+        ">",
+        "<",
+        "@",
+        "`",
+        ",",
+        "?",
+        "!",
+    ]
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def word_tokenize(self, word):
+        word = word.lower()
+        word = word.replace(",", "").replace("?", "").replace("'s", " 's")
+        return word.strip()
+
+    def process_punctuation(self, in_text):
+        out_text = in_text
+        for p in self.PUNCTUATIONS:
+            if (p + " " in in_text or " " + p in in_text) or (
+                re.search(self.COMMA_STRIP, in_text) is not None
+            ):
+                out_text = out_text.replace(p, "")
+            else:
+                out_text = out_text.replace(p, " ")
+        out_text = self.PERIOD_STRIP.sub("", out_text, re.UNICODE)
+        return out_text
+
+    def process_digit_article(self, in_text):
+        out_text = []
+        temp_text = in_text.lower().split()
+        for word in temp_text:
+            word = self.NUMBER_MAP.setdefault(word, word)
+            if word not in self.ARTICLES:
+                out_text.append(word)
+            else:
+                pass
+        for word_id, word in enumerate(out_text):
+            if word in self.CONTRACTIONS:
+                out_text[word_id] = self.CONTRACTIONS[word]
+        out_text = " ".join(out_text)
+        return out_text
+
+    def __call__(self, item):
+        item = self.word_tokenize(item)
+        item = item.replace("\n", " ").replace("\t", " ").strip()
+        item = self.process_punctuation(item)
+        item = self.process_digit_article(item)
+        return item
+
+
+class TextVQAAccuracyEvaluator:
+    def __init__(self):
+        self.answer_processor = EvalAIAnswerProcessor()
+
+    def _compute_answer_scores(self, raw_answers):
+        """
+        compute the accuracy (soft score) of human answers
+        """
+        answers = [self.answer_processor(a) for a in raw_answers]
+        #assert len(answers) == 10
+        gt_answers = list(enumerate(answers))
+        unique_answers = set(answers)
+        unique_answer_scores = {}
+
+        for unique_answer in unique_answers:
+            accs = []
+            for gt_answer in gt_answers:
+                other_answers = [item for item in gt_answers if item != gt_answer]
+                matching_answers = [
+                    item for item in other_answers if item[1] == unique_answer
+                ]
+                acc = min(1, float(len(matching_answers)) / 3)
+                accs.append(acc)
+            unique_answer_scores[unique_answer] = sum(accs) / len(accs)
+
+        return unique_answer_scores
+
+    def eval_pred_list(self, pred_list):
+        pred_scores = []
+        for entry in tqdm(pred_list):
+            pred_answer = self.answer_processor(entry["pred_answer"])
+            unique_answer_scores = self._compute_answer_scores(entry["gt_answers"])
+            score = unique_answer_scores.get(pred_answer, 0.0)
+            pred_scores.append(score)
+
+        accuracy = sum(pred_scores) / len(pred_scores)
+        return accuracy

QWENVL-PRIVATE/analysis/merge_jsonl.py

@@ -0,0 +1,32 @@
+import os
+
+# 1. Define your source files
+jsonl_file_list = [
+    '/home/yuheng/datasets/visual_cot_infovqa_subset15k.jsonl',
+    '/home/yuheng/datasets/visual_cot_docvqa_subset33k.jsonl',
+    '/home/yuheng/datasets/visual_cot_textvqa_subset18k.jsonl',
+    '/home/yuheng/datasets/chartqa_28k_qa.jsonl'
+]
+
+# 2. Define your destination
+dst_jsonl_path = '/home/yuheng/datasets/visual_cot_textual_subset94k.jsonl'
+
+print(f"Merging {len(jsonl_file_list)} files into {dst_jsonl_path}...")
+
+# 3. Merging Logic
+with open(dst_jsonl_path, 'w', encoding='utf-8') as outfile:
+    for file_path in jsonl_file_list:
+        print(f"Reading: {file_path}")
+        try:
+            with open(file_path, 'r', encoding='utf-8') as infile:
+                for line in infile:
+                    # stripping whitespace and re-adding \n ensures we don't
+                    # accidentally merge two JSON objects onto the same line
+                    # if a file was missing an EOF newline.
+                    clean_line = line.strip()
+                    if clean_line:
+                        outfile.write(clean_line + '\n')
+        except FileNotFoundError:
+            print(f"Warning: File not found {file_path}")
+
+print("Merge complete.")

QWENVL-PRIVATE/analysis/merge_pkls.py

@@ -0,0 +1,49 @@
+import pickle
+
+#only contains dict in the form of {question_id: item_dict} in each pkl
+pickle_file_list = [
+    '/root/epfs/sproject/datasets/stage3results/gated_vcot_docvqa33k_qwen2.5vl_3b.pkl',
+    '/root/epfs/sproject/datasets/stage3results/gated_chartqa28k_qwen2.5vl_3b.pkl',
+    '/root/epfs/sproject/datasets/stage3results/gated_vcot_llava68k_qwen2.5vl_3b.pkl'
+]
+
+dst_pkl_path = '/root/epfs/sproject/datasets/qwen2.5vl_3b_stage3.pkl'
+
+#pickle_file_list = [
+#     '/root/epfs/sproject/datasets/roi_vcot68k_ms.pkl',
+#     '/root/epfs/sproject/datasets/roi_vcot_docvqa33k_ms.pkl'
+# ]
+#
+# dst_pkl_path = '/root/epfs/sproject/datasets/roi_vcot101k_ms.pkl'
+
+
+merged_data = {}
+cur_data_idx = 0
+for pkl_path in pickle_file_list:
+    print(f"Loading data from {pkl_path}...")
+    try:
+        with open(pkl_path, 'rb') as f:
+            data = pickle.load(f)
+        if merged_data =={}:
+            merged_data = data
+            cur_data_idx = len(merged_data)
+        else:
+            for key, value in data.items():
+                if type(value['text']) is list:
+                    value['text'] = value['text'][0]
+                if key in merged_data:
+                    updated_key = cur_data_idx + key
+                    merged_data[updated_key] = value
+                else:
+                    merged_data[key] = value
+            cur_data_idx = len(merged_data)
+    except Exception as e:
+        print(f"Error loading pickle {pkl_path}: {e}")
+
+print(f"Saving merged data to {dst_pkl_path}...")
+try:
+    with open(dst_pkl_path, 'wb') as f:
+        pickle.dump(merged_data, f)
+    print("Merge and save completed successfully.")
+except Exception as e:
+    print(f"Error saving merged pickle: {e}")

QWENVL-PRIVATE/analysis/merge_stage1_pkls.py

@@ -0,0 +1,36 @@
+import os
+import pickle
+import argparse
+
+
+def merge_pkls(pkl_list, dst_path, dst_pkl_name='llava_v1_5_pseudo_roi.pkl'):
+    """
+    Get the pkl files from the list and save them to a single pkl file.
+    """
+    all_results = []
+    for pkl_path in pkl_list:
+        tmp_pkl_content = pickle.load(open(pkl_path, "rb"))
+        all_results.extend(tmp_pkl_content)
+    os.makedirs(dst_path, exist_ok=True)
+    dst_pkl_path = os.path.join(dst_path, dst_pkl_name)
+    with open(dst_pkl_path, "wb") as f:
+        pickle.dump(all_results, f)
+    print(f"Saved {len(all_results)} results to {dst_pkl_path}")
+    return dst_pkl_path
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Merge stage1 pkl files")
+    parser.add_argument("--pkl-files", nargs="+", required=True,
+                        help="List of pkl files to merge")
+    parser.add_argument("--dst-path", type=str, required=True,
+                        help="Destination directory for merged pkl")
+    parser.add_argument("--dst-pkl-name", type=str, required=True,
+                        help="Output pkl filename")
+    args = parser.parse_args()
+
+    merge_pkls(
+        pkl_list=args.pkl_files,
+        dst_path=args.dst_path,
+        dst_pkl_name=args.dst_pkl_name,
+    )

QWENVL-PRIVATE/analysis/plot_trade_off.py

@@ -0,0 +1,148 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+# --- 1. Data Definition ---
+# Benchmarks and their sample counts
+# Note: TextVQA uses 1250 samples as per user specification (val set subset or specific config)
+bench_keys = ['DocVQA', 'InfoVQA', 'TextVQA', 'V*', 'RW-Lite']
+samples = {
+    'DocVQA': 5349,
+    'InfoVQA': 701,
+    'TextVQA': 1250,
+    'V*': 191,
+    'RW-Lite': 960
+}
+
+# Raw Data: {Label: {'accs': [Doc, Info, Text, V*, RW], 'times': [Doc, Info, Text, V*, RW]}}
+# Data extracted from provided experimental results table
+
+raw_data_baseline = {
+    'Base\nMax-256': {
+        'accs': [81.2, 54.0, 75.5, 53.4, 27.6],
+        'times': [1330, 567, 938, 33, 292]
+    },
+    'Base\nMax-576': {
+        'accs': [92.0, 70.1, 81.1, 64.4, 35.4],
+        'times': [1990, 904, 1567, 56, 532]
+    },
+    'Base\nMax-1024': {
+        'accs': [93.9, 76.5, 82.8, 70.2, 37.5],
+        'times': [3066, 1423, 2471, 94, 896]
+    },
+    'Base\nMax-2048': {
+        'accs': [94.7, 80.4, 82.9, 72.3, 39.3],
+        'times': [6053, 2047, 2600, 213, 1925]
+    }
+}
+
+raw_data_ours = {
+    'Gated\nMax-256': {
+        'accs': [89.1, 69.6, 81.8, 67.5, 37.3],
+        # Note: RW-Lite time estimated as 435s based on avg ~1.5x slowdown from baseline (292s)
+        'times': [1919, 880, 1285, 53, 435]
+    },
+    'Gated\nMax-576': {
+        'accs': [93.8, 78.6, 82.8, 80.1, 43.4],
+        'times': [2430, 1277, 1878, 85, 813]
+    },
+    'Gated\nMax-1024': {
+        'accs': [94.6, 80.8, 83.9, 82.2, 46.6],
+        'times': [3481, 1801, 2677, 130, 1420]
+    }
+}
+
+
+# --- 2. Data Processing ---
+def calculate_metrics(raw_data):
+    avg_speeds = []
+    avg_accs = []
+    labels = []
+
+    for label, data in raw_data.items():
+        accs = data['accs']
+        times = data['times']
+
+        # Calculate samples/second for each benchmark
+        task_speeds = []
+        for i, task in enumerate(bench_keys):
+            n_samples = samples[task]
+            t_seconds = times[i]
+            if t_seconds > 0:
+                task_speeds.append(n_samples / t_seconds)
+            else:
+                task_speeds.append(0)
+
+        # Weighted Average Speed (Average of speeds)
+        avg_speed = np.mean(task_speeds)
+
+        # Average Accuracy
+        avg_acc = np.mean(accs)
+
+        avg_speeds.append(avg_speed)
+        avg_accs.append(avg_acc)
+        labels.append(label)
+
+    return avg_speeds, avg_accs, labels
+
+
+base_x, base_y, base_labels = calculate_metrics(raw_data_baseline)
+ours_x, ours_y, ours_labels = calculate_metrics(raw_data_ours)
+
+# --- 3. Plotting ---
+plt.figure(figsize=(10, 7))
+
+# Plot Baseline Curve
+plt.plot(base_x, base_y, marker='o', linestyle='--', color='#7f8c8d', label='Baseline', linewidth=2, markersize=9,
+         markerfacecolor='white', markeredgewidth=2)
+
+# Plot Ours Curve
+plt.plot(ours_x, ours_y, marker='*', linestyle='-', color='#e74c3c', label='Ours-Gated', linewidth=2, markersize=14,
+         markerfacecolor='#e74c3c')
+
+# Annotate Points (Baseline)
+for x, y, label in zip(base_x, base_y, base_labels):
+    offset = (0, -20)
+    if '2048' in label:
+        offset = (0, 10)
+        ha_align = 'right'
+    else:
+        ha_align = 'center'
+
+    plt.annotate(label, (x, y), xytext=offset, textcoords='offset points',
+                 ha=ha_align, fontsize=9, color='#555', weight='normal',
+                 bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.7))
+
+# Annotate Points (Ours)
+for x, y, label in zip(ours_x, ours_y, ours_labels):
+    offset = (0, 10)
+    if '1024' in label:
+        offset = (10, 0)
+        ha_align = 'left'
+    else:
+        ha_align = 'center'
+
+    plt.annotate(label, (x, y), xytext=offset, textcoords='offset points',
+                 ha=ha_align, fontsize=9, fontweight='bold', color='#c0392b',
+                 bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.7))
+
+# --- 4. Styling ---
+plt.title('Accuracy vs. Efficiency Trade-off\n(Average across DocVQA, InfoVQA, TextVQA, V*, RW-Lite)', fontsize=14,
+          pad=15)
+plt.xlabel('Average Throughput (Samples / Second) \n← Slower | Faster →', fontsize=12, fontweight='bold')
+plt.ylabel('Average Accuracy (%)', fontsize=12, fontweight='bold')
+plt.grid(True, linestyle=':', alpha=0.6)
+plt.legend(fontsize=11, loc='lower right')
+
+# Add Arrows for interpretation
+plt.annotate('', xy=(max(base_x) * 0.95, min(base_y) * 1.02), xytext=(max(base_x) * 0.85, min(base_y) * 1.02),
+             arrowprops=dict(arrowstyle="->", color='green', lw=1.5))
+plt.text(max(base_x) * 0.90, min(base_y) * 1.01, 'Higher Efficiency', color='green', ha='center', fontsize=9)
+
+plt.annotate('', xy=(min(base_x) * 1.05, max(ours_y) * 0.98), xytext=(min(base_x) * 1.05, max(ours_y) * 0.90),
+             arrowprops=dict(arrowstyle="->", color='green', lw=1.5))
+plt.text(min(base_x) * 1.06, max(ours_y) * 0.94, 'Higher Accuracy', color='green', rotation=90, va='center', fontsize=9)
+
+plt.tight_layout()
+plt.show()
+# plt.savefig('acc_efficiency_tradeoff_updated.png', dpi=300)
+# print("Plot saved as 'acc_efficiency_tradeoff_updated.png'")

QWENVL-PRIVATE/analysis/plot_vstar_tradeoff.py

@@ -0,0 +1,101 @@
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+import numpy as np
+
+# --- Aesthetic Setup (Apply principles from reference figures) ---
+# Use a serif font similar to LaTeX/academic papers
+mpl.rcParams['font.family'] = 'Times New Roman'
+mpl.rcParams['mathtext.fontset'] = 'stix'  # Ensures math text matches
+mpl.rcParams['axes.labelweight'] = 'bold'
+mpl.rcParams['axes.titleweight'] = 'bold'
+
+# --- Data Extraction (same as before) ---
+data = {
+    'Qwen2.5-VL-3B': {
+        'Baseline': {
+            'perf': [47.1, 52.9, 59.2, 65.4, 68.1, 70.2, 71.2, 73.3],
+            'throughput': [24.0, 17.4, 9.5, 5.7, 3.5, 2.2, 1.2, 0.8]
+        },
+        '+Ours': {
+            'perf': [53.9, 63.9, 70.7, 75.4, 76.4, 79.1, 78.5, 81.2],
+            'throughput': [16.1, 9.5, 5.2, 3.3, 2.0, 1.4, 1.0, 0.6]
+        }
+    },
+    'Qwen2.5-VL-7B': {
+        'Baseline': {
+            'perf': [42.9, 53.4, 64.4, 70.2, 72.8, 72.3, 74.9, 78.0],
+            'throughput': [21.1, 15.5, 9.1, 5.4, 3.4, 2.4, 1.4, 1.0]
+        },
+        '+Ours': {
+            'perf': [59.7, 70.7, 77.5, 85.9, 86.4, 83.2, 89.5, 87.9],
+            'throughput': [13.0, 8.6, 5.1, 3.0, 2.5, 2.1, 1.3, 0.9]
+        }
+    }
+}
+resolutions = [128, 256, 576, 1024, 1536, 2048, 3072, 4096]
+
+# --- Plotting ---
+fig, axs = plt.subplots(1, 2, figsize=(15, 5), sharey=False, dpi=300)
+
+# NEW: Professional color palette (colorblind-friendly)
+colors = {'Baseline': '#FF9900', '+Ours': '#00BFAE'}
+markers = {'Baseline': 'o', '+Ours': 's'}
+text_offsets = {'Baseline': (0, -2.5), '+Ours': (0, 2.5)}  # Adjusted offsets for new fonts/sizes
+
+model_names = ['Qwen2.5-VL-3B', 'Qwen2.5-VL-7B']
+
+for i, model_name in enumerate(model_names):
+    ax = axs[i]
+    for method in ['Baseline', '+Ours']:
+        perf = data[model_name][method]['perf']
+        throughput = data[model_name][method]['throughput']
+
+        # MODIFIED: Added zorder, markeredgecolor, and other style params
+        ax.plot(throughput, perf, marker=markers[method], color=colors[method],
+                linestyle='--', label=method, markersize=9,
+                markeredgewidth=1.0, markeredgecolor=colors[method],  # Add edge to markers
+                linewidth=2.0, zorder=3)  # zorder places plot lines above grid
+
+        # Annotate points with their resolution
+        for j, res in enumerate(resolutions):
+            # Only annotate key resolutions to avoid clutter
+            if res in [128, 256, 576, 1024, 1536, 2048, 3072, 4096]:
+                ax.text(throughput[j], perf[j] + text_offsets[method][1], str(res),
+                        fontsize=12, color=colors[method], ha='center', va='center', zorder=4, weight='bold')
+
+
+    # --- Major Aesthetic Overhaul for each subplot ---
+    ax.set_title(model_name, fontsize=18, pad=10)
+    ax.set_xscale('log')
+    from matplotlib.ticker import ScalarFormatter
+
+    ax.xaxis.set_major_formatter(ScalarFormatter())
+    ax.set_xticks([1, 2, 5, 10, 20])
+
+    ax.set_xlabel('Throughput (tokens/s)', fontsize=16)
+    if i == 0:
+        ax.set_ylabel('V* Performance', fontsize=16)
+
+    # NEW: Remove top and right spines for a cleaner look
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+
+    # NEW: Make remaining spines and ticks thicker and black
+    ax.spines['left'].set_linewidth(1.5)
+    ax.spines['bottom'].set_linewidth(1.5)
+    ax.spines['left'].set_color('black')
+    ax.spines['bottom'].set_color('black')
+    ax.tick_params(width=1.5, colors='black', direction='out', length=6)
+
+    # NEW: Make grid more subtle
+    ax.grid(True, which="both", ls="--", linewidth=0.7, color='lightgray', zorder=0)
+
+    # MODIFIED: Make legend cleaner (no frame)
+    ax.legend(prop={'weight': 'bold', 'size': 14}, loc='best', frameon=False)
+    ax.tick_params(axis='both', which='major', labelsize=14)
+
+# --- Final Touches ---
+# fig.suptitle('Performance-Throughput Trade-off on V* Benchmark', fontsize=22, y=0.95)
+plt.tight_layout(rect=[0, 0, 1, 0.96])
+plt.savefig('performance_throughput_tradeoff_final.pdf', bbox_inches='tight')
+plt.show()

QWENVL-PRIVATE/analysis/qwen3vl_debug.py

@@ -0,0 +1,410 @@
+import sys
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import AutoProcessor
+from qwen_src.qwen3_vl import Qwen3VLForConditionalGeneration
+import json
+import os
+from tqdm import tqdm
+from PIL import Image, ImageDraw, ImageFont
+from qwen_src.mm_utils import expand2square, create_pseudo_labels
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map, visualize_token_wise_attention
+import matplotlib.patches as patches
+from matplotlib.colors import ListedColormap, BoundaryNorm
+from analysis.denoise_textual_file_wdet import process_single_sample_packed
+PATCH_SIZE = 32 #32x32 patches For Qwen3-VL
+import pickle
+use_act = True
+import math
+
+def plot_pseudo_label(image_path, layer2_attn, grounding_attn_o2i, grounding_attn, is_sink_token = None,
+                            save_path=None, figsize=(15, 4), fg_bbox = None):
+    """
+    Plot image and attention maps.
+
+    Args:
+        image_path: str, path to the image file
+        layer2_attn: ndarray (576,), layer 2 attention weights
+        grounding_attn_o2i: ndarray (576,), grounding attention object to image
+        grounding_attn: ndarray (576,), grounding attention
+        is_sink_token: ndarray (576,) in bool type to indicate whether this token is a sink token
+        save_path: str, optional path to save the plot
+        figsize: tuple, figure size
+    """
+    # Create figure with subplots
+    fig, axes = plt.subplots(1, 4, figsize=figsize, dpi=200)
+
+    # Plot 1: Original Image
+    if os.path.exists(image_path):
+        img = Image.open(image_path)
+        axes[0].imshow(img)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+    else:
+        axes[0].text(0.5, 0.5, f'Image not found:\n{image_path}',
+                     ha='center', va='center', transform=axes[0].transAxes)
+        axes[0].set_title('Original Image')
+        axes[0].axis('off')
+
+    # Reshape attention arrays to 24x24
+    if layer2_attn.ndim == 1:
+        square_size_hw = int(math.sqrt(len(layer2_attn.flatten())))
+        square_size = (square_size_hw, square_size_hw)
+    else:
+        square_size = layer2_attn.shape
+    layer2_attn_2d = layer2_attn.reshape(square_size)
+    grounding_attn_o2i_2d = grounding_attn_o2i.reshape(square_size)
+    grounding_attn_2d = grounding_attn.reshape(square_size)
+    if is_sink_token is not None:
+        is_sink_token_2d = is_sink_token.reshape(square_size)
+
+
+    # Plot 2: Special handling for Pseudo-Label map vs. continuous attention
+    # Check if the data is likely discrete pseudo-labels (-1, 0, 1)
+    if -1 in layer2_attn:
+        # --- Discrete Colormap for Pseudo-Labels ---
+        axes[1].set_title('Pseudo Label Map')
+        # Define the colors: Ignore (-1)->Gray, Background (0)->Dark Purple, Foreground (1)->Yellow
+        colors = ['gray', 'black', '#FFE699']
+        cmap_discrete = ListedColormap(colors)
+        # Define the boundaries for the colormap: [-1.5, -0.5), [-0.5, 0.5), [0.5, 1.5]
+        boundaries = [-1.5, -0.5, 0.5, 1.5]
+        norm = BoundaryNorm(boundaries, cmap_discrete.N)
+        im1 = axes[1].imshow(layer2_attn_2d, cmap=cmap_discrete, norm=norm, aspect='auto')
+        # Customize the colorbar for discrete values
+        cbar = plt.colorbar(im1, ax=axes[1], shrink=0.6)
+        cbar.set_ticks([-1, 0, 1])
+        cbar.set_ticklabels(['Ignore', 'BG', 'FG'])
+    else:
+        # --- Continuous Colormap for standard attention ---
+        axes[1].set_title('Layer 2 Attention')
+        im1 = axes[1].imshow(layer2_attn_2d, cmap='viridis', aspect='auto')
+        plt.colorbar(im1, ax=axes[1], shrink=0.6)
+
+    # Plot 3: Grounding Attention O2I
+    im2 = axes[2].imshow(grounding_attn_o2i_2d, cmap='inferno', aspect='auto')
+    axes[2].set_title('Grounding Attn O2I')
+    axes[2].axis('off')
+    plt.colorbar(im2, ax=axes[2], shrink=0.6)
+
+    # Plot 4: Grounding Attention Q2I
+    im3 = axes[3].imshow(grounding_attn_2d, cmap='inferno', aspect='auto')
+    axes[3].set_title('Grounding Attn')
+    axes[3].axis('off')
+    plt.colorbar(im3, ax=axes[3], shrink=0.6)
+
+    # Highlight sink tokens on the attention maps
+    if is_sink_token is not None:
+        # Find the coordinates of all sink tokens
+        sink_coords = np.argwhere(is_sink_token_2d)
+        for row, col in sink_coords:
+            # Define the properties of the highlight box
+            rect_args = {
+                'xy': (col - 0.5, row - 0.5), # Bottom-left corner of the cell
+                'width': 1,
+                'height': 1,
+                'linewidth': 2.5,
+                'edgecolor': '#7CFC00',  # A bright, "LawnGreen" color
+                'facecolor': 'none'
+            }
+            # Add a new rectangle patch to each of the three attention plots
+            #axes[1].add_patch(patches.Rectangle(**rect_args))
+            axes[2].add_patch(patches.Rectangle(**rect_args))
+            #axes[3].add_patch(patches.Rectangle(**rect_args))
+
+    if fg_bbox is not None:
+        fg_rect_arg = {
+            'xy': (fg_bbox[1] - 0.5, fg_bbox[0] - 0.5),  # Bottom-left corner of the cell
+            'width': fg_bbox[3] - fg_bbox[1] + 1,
+            'height': fg_bbox[2] - fg_bbox[0] + 1,
+            'linewidth': 2.5,
+            'edgecolor': '#FF4500',  # A bright, "OrangeRed" color
+            'facecolor': 'none'
+        }
+        axes[1].add_patch(patches.Rectangle(**fg_rect_arg))
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+
+    plt.show()
+    pass
+
+def eval_model_qwen(model, processor, questions, image_folder,
+                    ocr_cache_file=None, is_natural_image=False, save_path=None, debug_samples=10):
+    
+    sample = 0
+    # questions = [q for q in questions if os.path.exists(os.path.join(image_folder, q['image']))]
+
+    ocr_cache = {}
+    if ocr_cache_file and os.path.exists(ocr_cache_file):
+        print(f"Loading existing OCR cache from {ocr_cache_file}...")
+        try:
+            with open(ocr_cache_file, 'rb') as f:
+                ocr_cache = pickle.load(f)
+            print(f"Loaded {len(ocr_cache)} cached entries.")
+        except Exception as e:
+            print(f"Failed to load cache: {e}. Starting fresh.")
+
+
+    for line in tqdm(questions[sample:]):
+        if use_act: activations.clear()
+        image_file = line["image"]
+        qs = line["text"] if 'text' in line else line['prompt']
+
+        if is_natural_image:
+            qs = qs.replace('Answer the question using a single word or phrase.',
+                            'Output grounding bounding box related to the question in JSON.')
+        else:
+            if 'Answer the question using a single word or phrase.' not in qs:
+                qs += ' Answer the question using a single word or phrase.'
+
+        image = Image.open(os.path.join(image_folder, image_file)).convert('RGB')
+        processor.image_processor.max_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        processor.image_processor.min_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        image_size = image.size
+        image = expand2square(image, (127, 127, 127))
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": f"{qs}"
+                    },
+                    # {
+                    #     "type": "image",
+                    #     "image": image#os.path.join(image_folder, image_file)
+                    # }, # 测试多图 done！
+                    {
+                        "type": "image",
+                        "image": image  # os.path.join(image_folder, image_file)
+                    }
+                ]
+            }
+        ]
+        # Preparation for inference
+        inputs = processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt"
+        )
+        inputs = inputs.to(model.device)
+
+        # Inference: Generation of the output
+        generated_ids = model.generate(**inputs, max_new_tokens=128, output_attentions=True)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+        print("Question:", qs)
+        print(output_text)
+
+        if use_act and activations.get('visual_mask', None) is None:
+            vision_start_token_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+            vision_end_token_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+            pos = inputs['input_ids'].tolist()[0].index(vision_start_token_id) + 1
+            pos_end = inputs['input_ids'].tolist()[0].index(vision_end_token_id)
+            activations['visual_mask'] = torch.zeros((1, inputs['input_ids'].shape[1]), dtype=torch.bool, device=device)
+            activations['visual_mask'][0, pos:pos_end] = True
+        output_shape = inputs['image_grid_thw'].cpu().numpy().squeeze(0)[1:] / 2
+        output_shape = np.array(output_shape, dtype=np.int32)
+        attn_weights_reduction = 'mean'  # 'none' or 'mean'
+        layers = model.config.text_config.num_hidden_layers
+        attn_weights_map = get_attn_weights_map(activations, layers=layers, reduction=attn_weights_reduction)
+
+
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', attn_size=output_shape, layers=layers, rgb_image=image,
+                                    #    save_path=f'{save_path}/{sample}_layer_attn.png')
+        # ###4B: 19, 23 ,24
+        # ##textual image heads
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                layer_idx=13, attn_size=output_shape,layers=layers, save_path=f'{save_path}/{sample}_13_layer_attn.png')#head: 24, 27
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                 layer_idx=14, attn_size=output_shape,layers=layers, save_path=f'{save_path}/{sample}_14_layer_attn.png')#head: 21, 23; sink_heads:6, 10
+        # ####natural image heads
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                layer_idx=14, attn_size=output_shape,layers=layers,save_path=f'{save_path}/{sample}_14_layer_attn.png')#head:
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                 layer_idx=17, attn_size=output_shape,layers=layers,save_path=f'{save_path}/{sample}_17_layer_attn.png')#head
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                layer_idx=25, attn_size=output_shape, layers=layers)  # head
+        ###8B
+        ###textual image heads
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                layer_idx=19, attn_size=output_shape,layers=layers)#head:
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                 layer_idx=24, attn_size=output_shape,layers=layers)#head
+        # ####natural image heads
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                layer_idx=20, attn_size=output_shape,layers=layers)#head:
+        # visualize_token_wise_attention(attn_weights_map, qk_type='output2images', avg_heads=False, rgb_image=image,
+        #                                 layer_idx=21, attn_size=output_shape,layers=layers)#head
+
+        required_heads = None
+
+        sink_heads = None
+
+        if "xxx" in args.model_path and use_act:
+            if is_natural_image: # gqa
+                required_heads = {14: [14, 15]}
+                sink_heads = {14: [3, 5]}
+            else: # textvqa
+                required_heads = {13: [4, 6, 7, 14, 15], 14:[6, 7, 14, 13, 15]}
+                sink_heads = {}
+
+        if "Qwen3-VL-4B" in args.model_path and use_act:
+            if is_natural_image: ##TODO: find heads for natural image
+                required_heads = {20: [15, 21, 23], 21: [8, 10, 11]}
+                sink_heads = {20: [1, 24]}
+            else:
+                ##textual image
+                required_heads = {19: [24, 27], 24: [21, 23]}
+                sink_heads = {24: [6, 10]}
+            required_layers = list(required_heads.keys())
+        elif "Qwen3-VL-8B" in args.model_path and use_act:
+            if is_natural_image:  ##TODO: find heads for natural image
+                required_heads = {20: [15, 21, 23], 21: [8, 10, 11]}
+                sink_heads = {21: [7]}
+            else:
+                ##textual image
+                required_heads = {19: [24, 27], 24: [21, 23]}
+                sink_heads = {24: [1, 2]}
+            required_layers = list(required_heads.keys())
+
+        if required_heads is None:
+            print(f"[WARN] No predefined required_heads for model '{args.model_path}'. "
+                  f"Only per-layer attention visualizations are saved. "
+                  f"Check {save_path}/{sample}_layer_attn.png to identify grounding heads.")
+            sample += 1
+            if sample >= debug_samples:
+                print('Debug samples limit reached. Stopping evaluation.')
+                break
+            continue
+
+        grounding_attn_o2i = np.zeros(output_shape)
+        total_grounding_heads = sum([len(heads) for heads in required_heads.values()])
+        for layer, heads in required_heads.items():
+            attn_weights = attn_weights_map[f'layer{layer}']['output2images'].reshape(-1, output_shape[0], output_shape[1])
+            for head in heads:
+                grounding_attn_o2i += attn_weights[head]
+        grounding_attn_o2i = grounding_attn_o2i / total_grounding_heads
+
+        sink_attn = np.zeros(output_shape)
+        total_sink_heads = sum([len(heads) for heads in sink_heads.values()])
+        if total_sink_heads > 0:
+            for layer, heads in sink_heads.items():
+                attn_weights = attn_weights_map[f'layer{layer}']['output2images'].reshape(-1, output_shape[0], output_shape[1])
+                for head in heads:
+                    sink_attn += attn_weights[head]
+        sink_attn = sink_attn / total_sink_heads
+
+        ##denoise with ocr boxes
+        img_path = os.path.join(image_folder, image_file)
+        # cached_boxes = ocr_cache.get(img_path, None)
+        # item = {
+        #     'image': img_path,
+        #     'grounding_attn_o2i': grounding_attn_o2i,
+        # }
+        # new_item, _, _ = process_single_sample_packed((item, None, cached_boxes))
+        # grounding_attn_o2i = new_item['grounding_attn_o2i']
+        if not is_natural_image:
+            cached_boxes = ocr_cache.get(img_path, None)
+            item = {
+                'image': img_path,
+                'grounding_attn_o2i': grounding_attn_o2i,
+            }
+            new_item, _, _ = process_single_sample_packed((item, None, cached_boxes))
+            grounding_attn_o2i = new_item['grounding_attn_o2i']
+
+        pseudo_set_denoised = create_pseudo_labels(sink_attn, grounding_attn_o2i, sink_thresh=0.02, bg_coff=0.05, binary_coff=0.25,
+                                          original_image_size=image_size, for_vis=True, mask_known_bg = True)
+        plot_pseudo_label(
+            image_path=img_path,
+            layer2_attn=pseudo_set_denoised['labels'],
+            grounding_attn_o2i=grounding_attn_o2i,
+            grounding_attn=pseudo_set_denoised['grounding_attn'],
+            is_sink_token=pseudo_set_denoised['sink_token_mask'],
+            fg_bbox=pseudo_set_denoised['fg_bbox'],
+            save_path=f'{save_path}/{sample}_pseudo_label.png',
+        )
+        sample += 1
+        if sample >= debug_samples:
+            print('Debug samples limit reached. Stopping evaluation.')
+            break
+        pass
+
+import argparse
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, required=True, help="model name")
+    parser.add_argument("--model_path", type=str, required=True, help="model path")
+    parser.add_argument(
+        "--dataset", type=str, required=True,
+        choices=["ocr_vqa", "DocVQA", "textvqa", "gqa"],
+        help="ocr_vqa, DocVQA, textvqa, gqa"
+    )
+    parser.add_argument("--save_path", type=str, required=True, help="save directory")
+
+    parser.add_argument("--questions_path", type=str, required=True, help="path to .jsonl/.json questions file")
+    parser.add_argument("--images_path", type=str, required=True, help="path to images root folder, not include images/")
+    parser.add_argument("--ocr_cache_file", type=str, default=None, help="path to ocr cache .pkl (optional)")
+
+    parser.add_argument("--debug_samples", type=int, default=100000, help="stop after N samples")
+    args = parser.parse_args()
+
+    device = "cuda"
+
+    model = Qwen3VLForConditionalGeneration.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+    ).eval().to(device)
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        padding_side="left",
+        use_fast=True,
+    )
+
+    if use_act:
+        activations, hooks = register_hooks(model, is_qwen3vl=True)
+
+    is_natural_image = (args.dataset == "gqa")
+
+    if args.dataset == "textvqa":
+        args.images_path = args.images_path + "/textvqa/train_images"
+
+    print("args.images_path: ", args.images_path)
+    ocr_cache_file = os.path.expanduser(args.ocr_cache_file) if args.ocr_cache_file else None
+
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.questions_path), "r")]
+    print('Total questions loaded:', len(questions))
+    questions = [q for q in questions if args.dataset in os.path.join(args.images_path, q["image"])]
+    print('Questions after filtering by target dataset:', len(questions))
+
+
+    eval_model_qwen(
+        model=model,
+        processor=processor,
+        questions=questions,
+        image_folder=args.images_path,
+        ocr_cache_file=ocr_cache_file,
+        is_natural_image=is_natural_image,
+        save_path=args.save_path,
+        debug_samples=args.debug_samples,
+    )
+
+    print(f"dataset: {args.dataset}")
+    print(f"questions_path: {args.questions_path}")
+    print(f"images_path: {args.images_path}")
+    print(f"ocr_cache_file: {ocr_cache_file}")
+    print(f"is_natural_image: {is_natural_image}")

QWENVL-PRIVATE/analysis/qwen3vl_debug_30b.py

@@ -0,0 +1,300 @@
+"""
+Generate ROI pseudo label maps for Qwen3-VL-30B-A3B (MoE).
+Same pipeline as qwen3vl_debug.py but adapted for:
+  1. Qwen3VLMoeForConditionalGeneration (MoE architecture)
+  2. device_map="auto" (58GB model, single 80GB GPU)
+  3. 48 layers (vs 36 for 4B/8B)
+
+Output: <save_path>/<sample_id>_pseudo_label.png (4-panel: image + pseudo + grounding_o2i + grounding)
+"""
+
+import sys
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
+import json
+import os
+from tqdm import tqdm
+from PIL import Image
+import math
+import pickle
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from qwen_src.mm_utils import expand2square, create_pseudo_labels
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map
+
+PATCH_SIZE = 32
+use_act = True
+
+
+def plot_pseudo_label(image_path, layer2_attn, grounding_attn_o2i, grounding_attn,
+                      is_sink_token=None, save_path=None, figsize=(15, 4), fg_bbox=None):
+    """Same visualization as qwen3vl_debug.py"""
+    from matplotlib.colors import ListedColormap, BoundaryNorm
+    import matplotlib.patches as patches
+
+    fig, axes = plt.subplots(1, 4, figsize=figsize, dpi=200)
+
+    if os.path.exists(image_path):
+        img = Image.open(image_path)
+        axes[0].imshow(img)
+    axes[0].set_title('Original Image')
+    axes[0].axis('off')
+
+    if layer2_attn.ndim == 1:
+        sq = int(math.sqrt(len(layer2_attn.flatten())))
+        square_size = (sq, sq)
+    else:
+        square_size = layer2_attn.shape
+    layer2_attn_2d = layer2_attn.reshape(square_size)
+    grounding_attn_o2i_2d = grounding_attn_o2i.reshape(square_size)
+    grounding_attn_2d = grounding_attn.reshape(square_size)
+
+    if -1 in layer2_attn:
+        colors = ['gray', 'black', '#FFE699']
+        cmap_discrete = ListedColormap(colors)
+        boundaries = [-1.5, -0.5, 0.5, 1.5]
+        norm = BoundaryNorm(boundaries, cmap_discrete.N)
+        im1 = axes[1].imshow(layer2_attn_2d, cmap=cmap_discrete, norm=norm, aspect='auto')
+        axes[1].set_title('Pseudo Label Map')
+        cbar = plt.colorbar(im1, ax=axes[1], shrink=0.6)
+        cbar.set_ticks([-1, 0, 1])
+        cbar.set_ticklabels(['Ignore', 'BG', 'FG'])
+    else:
+        im1 = axes[1].imshow(layer2_attn_2d, cmap='viridis', aspect='auto')
+        axes[1].set_title('Attention')
+        plt.colorbar(im1, ax=axes[1], shrink=0.6)
+
+    im2 = axes[2].imshow(grounding_attn_o2i_2d, cmap='inferno', aspect='auto')
+    axes[2].set_title('Grounding Attn O2I')
+    axes[2].axis('off')
+    plt.colorbar(im2, ax=axes[2], shrink=0.6)
+
+    im3 = axes[3].imshow(grounding_attn_2d, cmap='inferno', aspect='auto')
+    axes[3].set_title('Grounding Attn')
+    axes[3].axis('off')
+    plt.colorbar(im3, ax=axes[3], shrink=0.6)
+
+    if is_sink_token is not None:
+        is_sink_token_2d = is_sink_token.reshape(square_size)
+        sink_coords = np.argwhere(is_sink_token_2d)
+        for row, col in sink_coords:
+            rect_args = {'xy': (col - 0.5, row - 0.5), 'width': 1, 'height': 1,
+                         'linewidth': 2.5, 'edgecolor': '#7CFC00', 'facecolor': 'none'}
+            axes[2].add_patch(patches.Rectangle(**rect_args))
+
+    if fg_bbox is not None:
+        fg_rect_arg = {'xy': (fg_bbox[1] - 0.5, fg_bbox[0] - 0.5),
+                       'width': fg_bbox[3] - fg_bbox[1] + 1,
+                       'height': fg_bbox[2] - fg_bbox[0] + 1,
+                       'linewidth': 2.5, 'edgecolor': '#FF4500', 'facecolor': 'none'}
+        axes[1].add_patch(patches.Rectangle(**fg_rect_arg))
+
+    plt.tight_layout()
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+def eval_model_30b(model, processor, questions, image_folder,
+                   ocr_cache_file=None, is_natural_image=False, save_path=None, debug_samples=10):
+    sample = 0
+    ocr_cache = {}
+    if ocr_cache_file and os.path.exists(ocr_cache_file):
+        print(f"Loading OCR cache from {ocr_cache_file}...")
+        try:
+            ocr_cache = pickle.load(open(ocr_cache_file, 'rb'))
+            print(f"Loaded {len(ocr_cache)} cached entries.")
+        except Exception as e:
+            print(f"Failed to load cache: {e}")
+
+    device = next(model.parameters()).device  # get actual device from model
+
+    for line in tqdm(questions[sample:]):
+        if use_act:
+            activations.clear()
+        image_file = line["image"]
+        qs = line["text"] if 'text' in line else line['prompt']
+
+        if is_natural_image:
+            qs = qs.replace('Answer the question using a single word or phrase.',
+                            'Output grounding bounding box related to the question in JSON.')
+        else:
+            if 'Answer the question using a single word or phrase.' not in qs:
+                qs += ' Answer the question using a single word or phrase.'
+
+        img_path = os.path.join(image_folder, image_file)
+        if not os.path.exists(img_path):
+            continue
+
+        image = Image.open(img_path).convert('RGB')
+        processor.image_processor.max_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        processor.image_processor.min_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        image_size = image.size
+        image = expand2square(image, (127, 127, 127))
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"{qs}"},
+                {"type": "image", "image": image}
+            ]
+        }]
+        inputs = processor.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True,
+            return_dict=True, return_tensors="pt"
+        )
+        inputs = inputs.to(device)
+
+        with torch.no_grad():
+            generated_ids = model.generate(**inputs, max_new_tokens=128, output_attentions=True)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+        print(f"Q: {qs[:60]}...")
+        print(f"A: {output_text[0][:60]}...")
+
+        if use_act and activations.get('visual_mask', None) is None:
+            vision_start_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+            vision_end_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+            input_ids_list = inputs['input_ids'].tolist()[0]
+            pos = input_ids_list.index(vision_start_id) + 1
+            pos_end = input_ids_list.index(vision_end_id)
+            activations['visual_mask'] = torch.zeros(
+                (1, inputs['input_ids'].shape[1]), dtype=torch.bool, device=device)
+            activations['visual_mask'][0, pos:pos_end] = True
+
+        output_shape = inputs['image_grid_thw'].cpu().numpy().squeeze(0)[1:] / 2
+        output_shape = np.array(output_shape, dtype=np.int32)
+        layers = model.config.text_config.num_hidden_layers  # 48 for 30B
+        attn_weights_map = get_attn_weights_map(activations, layers=layers, reduction='mean')
+
+        # For 30B, no predefined heads yet. Use all-layer average as initial exploration.
+        # The user will identify grounding heads from these maps.
+        # Use layers 60-80% depth as initial guess (layers 28-38 for 48 layers)
+        required_heads = None
+
+        # First pass: save per-layer attention for head identification
+        # Use all layers average for grounding attention
+        grounding_attn_o2i = np.zeros(output_shape)
+        count = 0
+        for layer_idx in range(layers):
+            key = f'layer{layer_idx}'
+            if key in attn_weights_map and 'output2images' in attn_weights_map[key]:
+                attn = attn_weights_map[key]['output2images']
+                try:
+                    n_heads = attn.shape[0] // np.prod(output_shape)
+                    if n_heads > 0:
+                        attn_reshaped = attn.reshape(n_heads, *output_shape)
+                        grounding_attn_o2i += attn_reshaped.mean(axis=0)
+                        count += 1
+                except:
+                    pass
+        if count > 0:
+            grounding_attn_o2i /= count
+
+        # Sink attention: use first few layers
+        sink_attn = np.zeros(output_shape)
+        sink_count = 0
+        for layer_idx in range(min(5, layers)):
+            key = f'layer{layer_idx}'
+            if key in attn_weights_map and 'output2images' in attn_weights_map[key]:
+                attn = attn_weights_map[key]['output2images']
+                try:
+                    n_heads = attn.shape[0] // np.prod(output_shape)
+                    if n_heads > 0:
+                        attn_reshaped = attn.reshape(n_heads, *output_shape)
+                        sink_attn += attn_reshaped.mean(axis=0)
+                        sink_count += 1
+                except:
+                    pass
+        if sink_count > 0:
+            sink_attn /= sink_count
+
+        # Denoise with OCR boxes (for textual images)
+        if not is_natural_image:
+            try:
+                from analysis.denoise_textual_file_wdet import process_single_sample_packed
+                cached_boxes = ocr_cache.get(img_path, None)
+                item = {'image': img_path, 'grounding_attn_o2i': grounding_attn_o2i}
+                new_item, _, _ = process_single_sample_packed((item, None, cached_boxes))
+                grounding_attn_o2i = new_item['grounding_attn_o2i']
+            except Exception as e:
+                print(f"  Denoise skipped: {e}")
+
+        # Create pseudo labels
+        pseudo_set = create_pseudo_labels(
+            sink_attn, grounding_attn_o2i,
+            sink_thresh=0.02, bg_coff=0.05, binary_coff=0.25,
+            original_image_size=image_size, for_vis=True, mask_known_bg=True
+        )
+
+        plot_pseudo_label(
+            image_path=img_path,
+            layer2_attn=pseudo_set['labels'],
+            grounding_attn_o2i=grounding_attn_o2i,
+            grounding_attn=pseudo_set['grounding_attn'],
+            is_sink_token=pseudo_set['sink_token_mask'],
+            fg_bbox=pseudo_set['fg_bbox'],
+            save_path=f'{save_path}/{sample}_pseudo_label.png',
+        )
+        sample += 1
+        if sample >= debug_samples:
+            print('Debug samples limit reached.')
+            break
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, default="Qwen3-VL-30B-A3B-Instruct")
+    parser.add_argument("--model_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/model/Qwen3-VL-30B-A3B-Instruct")
+    parser.add_argument("--dataset", type=str, default="gqa", choices=["gqa", "textvqa", "ocr_vqa"])
+    parser.add_argument("--save_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/stage0/Qwen3-VL-30B-A3B-Instruct/gqa")
+    parser.add_argument("--questions_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/llava_v1_5_mix665k_selected_qa.jsonl")
+    parser.add_argument("--images_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/")
+    parser.add_argument("--ocr_cache_file", type=str, default=None)
+    parser.add_argument("--debug_samples", type=int, default=10)
+    args = parser.parse_args()
+
+    os.makedirs(args.save_path, exist_ok=True)
+    device = "cuda"
+
+    print(f"Loading {args.model_name} (MoE: 128 experts, 48 layers)...")
+    model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+        device_map="auto",
+    ).eval()
+    print(f"Model loaded on device_map: {getattr(model, 'hf_device_map', 'unknown')}")
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_path, trust_remote_code=True, padding_side="left", use_fast=True,
+    )
+
+    activations, hooks = register_hooks(model, is_qwen3vl=True)
+
+    is_natural_image = (args.dataset == "gqa")
+    if args.dataset == "textvqa":
+        args.images_path = args.images_path + "/textvqa/train_images"
+
+    questions = [json.loads(q) for q in open(args.questions_path)]
+    questions = [q for q in questions if args.dataset in os.path.join(args.images_path, q["image"])]
+    print(f"GQA questions: {len(questions)}")
+
+    eval_model_30b(
+        model=model, processor=processor, questions=questions,
+        image_folder=args.images_path, ocr_cache_file=args.ocr_cache_file,
+        is_natural_image=is_natural_image, save_path=args.save_path,
+        debug_samples=args.debug_samples,
+    )
+    for h in hooks:
+        h.remove()

QWENVL-PRIVATE/analysis/qwen3vl_debug_32b.py

@@ -0,0 +1,349 @@
+"""
+Generate ROI pseudo label maps for Qwen3-VL-32B-Instruct (DENSE).
+
+32B is the dense variant (model_type=qwen3_vl), NOT MoE — so the pipeline
+is fully aligned with 4B/8B dense (qwen3vl_debug.py). Only differences:
+  1. 64 layers (vs 36 for 4B/8B, 48 for 30B MoE)
+  2. 64 attention heads (vs varying for 4B/8B)
+  3. hidden_size=5120 (vs 2560 for 4B, 3584 for 8B)
+  4. ~64GB in bf16 → device_map="auto" required (won't fit on single 80GB GPU
+     when accounting for activations + KV cache during generation)
+  5. DeepStack indexes [8, 16, 24] (same as 30B/8B)
+
+Required grounding/sink heads for 32B are not yet identified — first pass
+saves per-layer attention maps (set DEBUG_LAYER_VIS=1) so the user can
+inspect and pick heads. Second pass uses the identified heads.
+
+Output: <save_path>/<sample_id>_pseudo_label.png (4-panel figure)
+"""
+
+import sys
+import os
+import math
+import json
+import pickle
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+from PIL import Image
+from transformers import AutoProcessor
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from qwen_src.qwen3_vl import Qwen3VLForConditionalGeneration
+from qwen_src.mm_utils import expand2square, create_pseudo_labels
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map, visualize_token_wise_attention
+
+PATCH_SIZE = 32  # 32x32 patches for Qwen3-VL
+use_act = True
+
+
+def plot_pseudo_label(image_path, layer2_attn, grounding_attn_o2i, grounding_attn,
+                      is_sink_token=None, save_path=None, figsize=(15, 4), fg_bbox=None):
+    """4-panel figure: original image | pseudo label | grounding o2i | grounding."""
+    from matplotlib.colors import ListedColormap, BoundaryNorm
+    import matplotlib.patches as patches
+
+    fig, axes = plt.subplots(1, 4, figsize=figsize, dpi=200)
+
+    if os.path.exists(image_path):
+        img = Image.open(image_path)
+        axes[0].imshow(img)
+    axes[0].set_title('Original Image')
+    axes[0].axis('off')
+
+    if layer2_attn.ndim == 1:
+        sq = int(math.sqrt(len(layer2_attn.flatten())))
+        square_size = (sq, sq)
+    else:
+        square_size = layer2_attn.shape
+    layer2_attn_2d = layer2_attn.reshape(square_size)
+    grounding_attn_o2i_2d = grounding_attn_o2i.reshape(square_size)
+    grounding_attn_2d = grounding_attn.reshape(square_size)
+
+    if -1 in layer2_attn:
+        colors = ['gray', 'black', '#FFE699']
+        cmap_discrete = ListedColormap(colors)
+        boundaries = [-1.5, -0.5, 0.5, 1.5]
+        norm = BoundaryNorm(boundaries, cmap_discrete.N)
+        im1 = axes[1].imshow(layer2_attn_2d, cmap=cmap_discrete, norm=norm, aspect='auto')
+        axes[1].set_title('Pseudo Label Map')
+        cbar = plt.colorbar(im1, ax=axes[1], shrink=0.6)
+        cbar.set_ticks([-1, 0, 1])
+        cbar.set_ticklabels(['Ignore', 'BG', 'FG'])
+    else:
+        im1 = axes[1].imshow(layer2_attn_2d, cmap='viridis', aspect='auto')
+        axes[1].set_title('Attention')
+        plt.colorbar(im1, ax=axes[1], shrink=0.6)
+
+    im2 = axes[2].imshow(grounding_attn_o2i_2d, cmap='inferno', aspect='auto')
+    axes[2].set_title('Grounding Attn O2I')
+    axes[2].axis('off')
+    plt.colorbar(im2, ax=axes[2], shrink=0.6)
+
+    im3 = axes[3].imshow(grounding_attn_2d, cmap='inferno', aspect='auto')
+    axes[3].set_title('Grounding Attn')
+    axes[3].axis('off')
+    plt.colorbar(im3, ax=axes[3], shrink=0.6)
+
+    if is_sink_token is not None:
+        is_sink_token_2d = is_sink_token.reshape(square_size)
+        sink_coords = np.argwhere(is_sink_token_2d)
+        for row, col in sink_coords:
+            rect_args = {'xy': (col - 0.5, row - 0.5), 'width': 1, 'height': 1,
+                         'linewidth': 2.5, 'edgecolor': '#7CFC00', 'facecolor': 'none'}
+            axes[2].add_patch(patches.Rectangle(**rect_args))
+
+    if fg_bbox is not None:
+        fg_rect_arg = {'xy': (fg_bbox[1] - 0.5, fg_bbox[0] - 0.5),
+                       'width': fg_bbox[3] - fg_bbox[1] + 1,
+                       'height': fg_bbox[2] - fg_bbox[0] + 1,
+                       'linewidth': 2.5, 'edgecolor': '#FF4500', 'facecolor': 'none'}
+        axes[1].add_patch(patches.Rectangle(**fg_rect_arg))
+
+    plt.tight_layout()
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+def eval_model_32b(model, processor, questions, image_folder,
+                   ocr_cache_file=None, is_natural_image=False, save_path=None,
+                   debug_samples=10, save_per_layer_vis=False):
+    sample = 0
+    ocr_cache = {}
+    if ocr_cache_file and os.path.exists(ocr_cache_file):
+        print(f"Loading OCR cache from {ocr_cache_file}...")
+        try:
+            ocr_cache = pickle.load(open(ocr_cache_file, 'rb'))
+            print(f"Loaded {len(ocr_cache)} cached entries.")
+        except Exception as e:
+            print(f"Failed to load cache: {e}")
+
+    # On multi-GPU device_map="auto", model.device may not exist. Use first param.
+    device = next(model.parameters()).device
+
+    for line in tqdm(questions[sample:]):
+        if use_act:
+            activations.clear()
+        image_file = line["image"]
+        qs = line["text"] if 'text' in line else line['prompt']
+
+        if is_natural_image:
+            qs = qs.replace('Answer the question using a single word or phrase.',
+                            'Output grounding bounding box related to the question in JSON.')
+        else:
+            if 'Answer the question using a single word or phrase.' not in qs:
+                qs += ' Answer the question using a single word or phrase.'
+
+        img_path = os.path.join(image_folder, image_file)
+        if not os.path.exists(img_path):
+            continue
+
+        image = Image.open(img_path).convert('RGB')
+        processor.image_processor.max_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        processor.image_processor.min_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        image_size = image.size
+        image = expand2square(image, (127, 127, 127))
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"{qs}"},
+                {"type": "image", "image": image}
+            ]
+        }]
+        inputs = processor.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True,
+            return_dict=True, return_tensors="pt"
+        )
+        inputs = inputs.to(device)
+
+        with torch.no_grad():
+            generated_ids = model.generate(**inputs, max_new_tokens=128, output_attentions=True)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+        print(f"Q: {qs[:80]}")
+        print(f"A: {output_text[0][:80]}")
+
+        if use_act and activations.get('visual_mask', None) is None:
+            vision_start_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+            vision_end_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+            input_ids_list = inputs['input_ids'].tolist()[0]
+            pos = input_ids_list.index(vision_start_id) + 1
+            pos_end = input_ids_list.index(vision_end_id)
+            activations['visual_mask'] = torch.zeros(
+                (1, inputs['input_ids'].shape[1]), dtype=torch.bool, device=device)
+            activations['visual_mask'][0, pos:pos_end] = True
+
+        output_shape = inputs['image_grid_thw'].cpu().numpy().squeeze(0)[1:] / 2
+        output_shape = np.array(output_shape, dtype=np.int32)
+        layers = model.config.text_config.num_hidden_layers  # 64 for 32B
+        attn_weights_map = get_attn_weights_map(activations, layers=layers, reduction='mean')
+
+        # Optional: dump per-layer attention maps for head identification.
+        # Set save_per_layer_vis=True (or env DEBUG_LAYER_VIS=1) on the first
+        # exploratory run, then update required_heads/sink_heads below from
+        # the visualizations and rerun with save_per_layer_vis=False.
+        if save_per_layer_vis:
+            visualize_token_wise_attention(
+                attn_weights_map, qk_type='output2images', attn_size=output_shape,
+                layers=layers, rgb_image=image,
+                save_path=f'{save_path}/{sample}_layer_attn.png',
+            )
+
+        # Required heads for 32B — TBD. Run with --save_per_layer_vis to identify
+        # the strongest grounding heads from the per-layer visualization, then
+        # populate the dicts below. For 4B/8B/30B-MoE these are hard-coded after
+        # exploration. For 32B (64 layers, 64 heads), use mean-of-mid-layers as
+        # a sensible default until proper heads are identified.
+        required_heads = None
+        sink_heads = None
+
+        if "Qwen3-VL-32B" in args.model_path and use_act:
+            if is_natural_image:
+                # TODO: identify natural-image grounding heads for 32B
+                # 32B mid-depth layers (~64% depth) are good candidates;
+                # use 38-40 as exploratory default until properly identified.
+                required_heads = {38: list(range(64)), 39: list(range(64)), 40: list(range(64))}
+                sink_heads = {}
+            else:
+                # TODO: identify textual grounding heads for 32B
+                # Mirror 8B structure (mid-late depth) as starting point;
+                # adjust after running with --save_per_layer_vis.
+                required_heads = {38: list(range(64)), 39: list(range(64)), 40: list(range(64))}
+                sink_heads = {}
+
+        if required_heads is None:
+            print(f"[WARN] No predefined required_heads for model '{args.model_path}'. "
+                  f"Per-layer attention visualizations saved if --save_per_layer_vis. "
+                  f"Inspect them at {save_path}/{sample}_layer_attn.png to pick grounding heads.")
+            sample += 1
+            if sample >= debug_samples:
+                print('Debug samples limit reached.')
+                break
+            continue
+
+        grounding_attn_o2i = np.zeros(output_shape)
+        total_grounding_heads = sum(len(h) for h in required_heads.values())
+        for layer, heads in required_heads.items():
+            attn = attn_weights_map[f'layer{layer}']['output2images'].reshape(-1, output_shape[0], output_shape[1])
+            for head in heads:
+                grounding_attn_o2i += attn[head]
+        grounding_attn_o2i = grounding_attn_o2i / total_grounding_heads
+
+        sink_attn = np.zeros(output_shape)
+        total_sink_heads = sum(len(h) for h in sink_heads.values()) if sink_heads else 0
+        if total_sink_heads > 0:
+            for layer, heads in sink_heads.items():
+                attn = attn_weights_map[f'layer{layer}']['output2images'].reshape(-1, output_shape[0], output_shape[1])
+                for head in heads:
+                    sink_attn += attn[head]
+            sink_attn = sink_attn / total_sink_heads
+
+        # Denoise with OCR boxes for textual images
+        if not is_natural_image:
+            try:
+                from analysis.denoise_textual_file_wdet import process_single_sample_packed
+                cached_boxes = ocr_cache.get(img_path, None)
+                item = {'image': img_path, 'grounding_attn_o2i': grounding_attn_o2i}
+                new_item, _, _ = process_single_sample_packed((item, None, cached_boxes))
+                grounding_attn_o2i = new_item['grounding_attn_o2i']
+            except Exception as e:
+                print(f"  Denoise skipped: {e}")
+
+        pseudo_set = create_pseudo_labels(
+            sink_attn, grounding_attn_o2i,
+            sink_thresh=0.02, bg_coff=0.05, binary_coff=0.25,
+            original_image_size=image_size, for_vis=True, mask_known_bg=True,
+        )
+        plot_pseudo_label(
+            image_path=img_path,
+            layer2_attn=pseudo_set['labels'],
+            grounding_attn_o2i=grounding_attn_o2i,
+            grounding_attn=pseudo_set['grounding_attn'],
+            is_sink_token=pseudo_set['sink_token_mask'],
+            fg_bbox=pseudo_set['fg_bbox'],
+            save_path=f'{save_path}/{sample}_pseudo_label.png',
+        )
+        sample += 1
+        if sample >= debug_samples:
+            print('Debug samples limit reached.')
+            break
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, default="Qwen3-VL-32B-Instruct")
+    parser.add_argument("--model_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/model/Qwen3-VL-32B-Instruct")
+    parser.add_argument("--dataset", type=str, default="gqa",
+                        choices=["ocr_vqa", "DocVQA", "textvqa", "gqa"])
+    parser.add_argument("--save_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/stage0/Qwen3-VL-32B-Instruct/gqa")
+    parser.add_argument("--questions_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/llava_v1_5_mix665k_selected_qa.jsonl")
+    parser.add_argument("--images_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/")
+    parser.add_argument("--ocr_cache_file", type=str, default=None)
+    parser.add_argument("--debug_samples", type=int, default=10)
+    parser.add_argument("--save_per_layer_vis", action="store_true",
+                        help="Save per-layer attention maps for head identification (slow).")
+    args = parser.parse_args()
+
+    os.makedirs(args.save_path, exist_ok=True)
+
+    print(f"Loading {args.model_name} (DENSE: 64 layers, hidden=5120, ~64GB bf16)...")
+    # device_map="auto" — 64GB model needs 2+ GPUs.
+    model = Qwen3VLForConditionalGeneration.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+        device_map="auto",
+    ).eval()
+    print(f"Model loaded on device_map: {getattr(model, 'hf_device_map', 'unknown')}")
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_path, trust_remote_code=True, padding_side="left", use_fast=True,
+    )
+
+    if use_act:
+        activations, hooks = register_hooks(model, is_qwen3vl=True)
+
+    is_natural_image = (args.dataset == "gqa")
+    if args.dataset == "textvqa":
+        args.images_path = args.images_path + "/textvqa/train_images"
+
+    # Allow override via env var as well
+    save_per_layer_vis = args.save_per_layer_vis or os.environ.get("DEBUG_LAYER_VIS", "0") == "1"
+
+    print("args.images_path:", args.images_path)
+    ocr_cache_file = os.path.expanduser(args.ocr_cache_file) if args.ocr_cache_file else None
+
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.questions_path), "r")]
+    print('Total questions loaded:', len(questions))
+    questions = [q for q in questions if args.dataset in os.path.join(args.images_path, q["image"])]
+    print('Questions after filtering by target dataset:', len(questions))
+
+    eval_model_32b(
+        model=model, processor=processor, questions=questions,
+        image_folder=args.images_path, ocr_cache_file=ocr_cache_file,
+        is_natural_image=is_natural_image, save_path=args.save_path,
+        debug_samples=args.debug_samples, save_per_layer_vis=save_per_layer_vis,
+    )
+
+    if use_act:
+        for h in hooks:
+            h.remove()
+
+    print(f"dataset: {args.dataset}")
+    print(f"questions_path: {args.questions_path}")
+    print(f"images_path: {args.images_path}")
+    print(f"save_path: {args.save_path}")
+    print(f"is_natural_image: {is_natural_image}")

QWENVL-PRIVATE/analysis/qwen3vl_identify_heads_30b.py

@@ -0,0 +1,202 @@
+"""
+Identify grounding/sink attention heads for Qwen3-VL-30B-A3B (MoE).
+Mirrors qwen3vl_debug.py but:
+  1. Uses Qwen3VLMoeForConditionalGeneration (MoE architecture)
+  2. Uses device_map="auto" (58GB model on 80GB GPU)
+  3. Generates per-layer attention maps for ALL 48 layers (head identification)
+  4. Also generates per-head maps for selected layers
+
+Output per sample:
+  <id>_layer_attn.png         — all layers averaged over heads (identify which layers matter)
+  <id>_layer<L>_heads.png     — all heads for layer L (identify which heads within a layer)
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python analysis/qwen3vl_identify_heads_30b.py \
+        --dataset gqa --debug_samples 10
+"""
+
+import sys
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
+import json
+import os
+from tqdm import tqdm
+from PIL import Image
+import math
+import pickle
+import argparse
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from qwen_src.mm_utils import expand2square
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map, visualize_token_wise_attention
+
+PATCH_SIZE = 32
+use_act = True
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, default="Qwen3-VL-30B-A3B-Instruct")
+    parser.add_argument("--model_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/model/Qwen3-VL-30B-A3B-Instruct")
+    parser.add_argument("--dataset", type=str, default="gqa", choices=["gqa", "textvqa", "ocr_vqa"])
+    parser.add_argument("--save_path", type=str, default=None)
+    parser.add_argument("--questions_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/llava_v1_5_mix665k_selected_qa.jsonl")
+    parser.add_argument("--images_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/")
+    parser.add_argument("--ocr_cache_file", type=str, default=None)
+    parser.add_argument("--debug_samples", type=int, default=10)
+    parser.add_argument("--inspect_layers", type=str, default=None,
+                        help="Comma-separated layer indices for per-head visualization, e.g. '24,30,36'")
+    args = parser.parse_args()
+
+    if args.save_path is None:
+        args.save_path = f"/mnt/bn/leonworkspace/terry/data/stage0/{args.model_name}/{args.dataset}"
+    os.makedirs(args.save_path, exist_ok=True)
+
+    device = "cuda"
+    is_natural_image = (args.dataset == "gqa")
+
+    print(f"Loading {args.model_name} (MoE, 48 layers, 32 heads)...")
+    model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+        device_map="auto",
+    ).eval()
+    print(f"Model loaded.")
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_path, trust_remote_code=True, padding_side="left", use_fast=True,
+    )
+
+    activations, hooks = register_hooks(model, is_qwen3vl=True)
+
+    layers = model.config.text_config.num_hidden_layers  # 48
+
+    # Parse inspect_layers
+    inspect_layers = None
+    if args.inspect_layers:
+        inspect_layers = [int(x) for x in args.inspect_layers.split(',')]
+        print(f"Will generate per-head maps for layers: {inspect_layers}")
+
+    # Adjust image path for textvqa
+    images_path = args.images_path
+    if args.dataset == "textvqa":
+        images_path = images_path + "/textvqa/train_images"
+
+    # Load questions
+    questions = [json.loads(q) for q in open(args.questions_path)]
+    questions = [q for q in questions if args.dataset in os.path.join(images_path, q["image"])]
+    print(f"{args.dataset} questions: {len(questions)}")
+
+    sample = 0
+    for line in tqdm(questions):
+        activations.clear()
+
+        image_file = line["image"]
+        qs = line["text"] if "text" in line else line["prompt"]
+
+        if is_natural_image:
+            qs = qs.replace('Answer the question using a single word or phrase.',
+                            'Output grounding bounding box related to the question in JSON.')
+        else:
+            if 'Answer the question using a single word or phrase.' not in qs:
+                qs += ' Answer the question using a single word or phrase.'
+
+        img_path = os.path.join(images_path, image_file)
+        if not os.path.exists(img_path):
+            continue
+
+        image = Image.open(img_path).convert('RGB')
+        processor.image_processor.max_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        processor.image_processor.min_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        image = expand2square(image, (127, 127, 127))
+
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": qs},
+                {"type": "image", "image": image},
+            ]
+        }]
+
+        inputs = processor.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True,
+            return_dict=True, return_tensors="pt"
+        )
+        inputs = inputs.to(device)
+
+        with torch.no_grad():
+            generated_ids = model.generate(**inputs, max_new_tokens=128, output_attentions=True)
+
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        answer = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        print(f"  [{sample}] Q: {qs[:60]}...")
+        print(f"       A: {answer[:60]}...")
+
+        # Build visual mask
+        if activations.get('visual_mask', None) is None:
+            vision_start_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+            vision_end_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+            input_ids_list = inputs['input_ids'].tolist()[0]
+            pos = input_ids_list.index(vision_start_id) + 1
+            pos_end = input_ids_list.index(vision_end_id)
+            activations['visual_mask'] = torch.zeros(
+                (1, inputs['input_ids'].shape[1]), dtype=torch.bool, device=device)
+            activations['visual_mask'][0, pos:pos_end] = True
+
+        output_shape = inputs['image_grid_thw'].cpu().numpy().squeeze(0)[1:] / 2
+        output_shape = np.array(output_shape, dtype=np.int32)
+
+        attn_weights_map = get_attn_weights_map(activations, layers=layers, reduction='mean')
+
+        # 1. Per-layer average attention map (all layers in one grid)
+        visualize_token_wise_attention(
+            attn_weights_map,
+            qk_type='output2images',
+            attn_size=output_shape,
+            layers=layers,
+            rgb_image=image,
+            save_path=f'{args.save_path}/{sample}_layer_attn.png'
+        )
+        print(f"       Saved: {sample}_layer_attn.png")
+
+        # 2. Per-head attention for specified layers
+        if inspect_layers:
+            for layer_idx in inspect_layers:
+                if layer_idx >= layers:
+                    continue
+                visualize_token_wise_attention(
+                    attn_weights_map,
+                    qk_type='output2images',
+                    avg_heads=False,
+                    rgb_image=image,
+                    layer_idx=layer_idx,
+                    attn_size=output_shape,
+                    layers=layers,
+                    save_path=f'{args.save_path}/{sample}_layer{layer_idx}_heads.png'
+                )
+                print(f"       Saved: {sample}_layer{layer_idx}_heads.png")
+
+        sample += 1
+        if sample >= args.debug_samples:
+            print('Debug samples limit reached.')
+            break
+
+    for h in hooks:
+        h.remove()
+    print(f"\nDone! {sample} samples saved to {args.save_path}/")
+    print(f"Next step: inspect *_layer_attn.png to find grounding layers,")
+    print(f"then re-run with --inspect_layers <L1,L2,...> for per-head maps.")
+
+
+if __name__ == "__main__":
+    main()

QWENVL-PRIVATE/analysis/qwen3vl_identify_heads_32b.py

@@ -0,0 +1,237 @@
+"""
+Identify grounding/sink attention heads for Qwen3-VL-32B-Instruct (DENSE).
+
+32B is the dense variant (model_type=qwen3_vl), NOT MoE — so the pipeline
+mirrors qwen3vl_debug.py exactly (same code path as 4B/8B dense). Relative to
+qwen3vl_identify_heads_30b.py:
+  1. Uses Qwen3VLForConditionalGeneration (dense, not MoE)
+  2. Uses device_map="auto" (64GB model — won't fit on single 80GB with
+     activations + KV cache, so shard across all visible GPUs)
+  3. 64 hidden layers, 64 attention heads, 8 KV heads (GQA 8:1),
+     hidden_size=5120 — auto-detected from config
+  4. Generates per-layer attention maps for ALL 64 layers
+  5. Also generates per-head maps for selected layers via --inspect_layers
+
+Output per sample:
+  <id>_layer_attn.png         — all 64 layers averaged over heads (find grounding layers)
+  <id>_layer<L>_heads.png     — all 64 heads for layer L (find grounding heads within)
+
+Usage (single-GPU is fine — device_map=auto handles sharding if multiple visible):
+    CUDA_VISIBLE_DEVICES=0,1 python analysis/qwen3vl_identify_heads_32b.py \\
+        --dataset gqa --debug_samples 10
+
+    # After first pass, re-run with chosen layers for per-head maps:
+    CUDA_VISIBLE_DEVICES=0,1 python analysis/qwen3vl_identify_heads_32b.py \\
+        --dataset gqa --debug_samples 10 --inspect_layers 38,40,42,44
+
+Next step after this: plug the identified grounding/sink heads into
+qwen3vl_debug_32b.py (required_heads / sink_heads dicts).
+"""
+
+import sys
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import json
+import os
+from tqdm import tqdm
+from PIL import Image
+import argparse
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from qwen_src.qwen3_vl import Qwen3VLForConditionalGeneration
+from qwen_src.mm_utils import expand2square
+from qwen_src.ana_utils import register_hooks, get_attn_weights_map, visualize_token_wise_attention
+from transformers import AutoProcessor
+
+PATCH_SIZE = 32
+use_act = True
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, default="Qwen3-VL-32B-Instruct")
+    parser.add_argument("--model_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/model/Qwen3-VL-32B-Instruct")
+    parser.add_argument("--dataset", type=str, default="gqa",
+                        choices=["gqa", "textvqa", "ocr_vqa", "gqa_cot", "textvqa_cot", "docvqa_cot"])
+    parser.add_argument("--save_path", type=str, default=None)
+    parser.add_argument("--questions_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/llava_v1_5_mix665k_selected_qa.jsonl")
+    parser.add_argument("--images_path", type=str,
+                        default="/mnt/bn/leonworkspace/terry/data/")
+    parser.add_argument("--ocr_cache_file", type=str, default=None)
+    parser.add_argument("--debug_samples", type=int, default=10)
+    parser.add_argument("--inspect_layers", type=str, default=None,
+                        help="Comma-separated layer indices for per-head visualization, e.g. '38,40,42,44'")
+    args = parser.parse_args()
+
+    if args.save_path is None:
+        args.save_path = f"/mnt/bn/leonworkspace/terry/data/stage0/{args.model_name}/{args.dataset}"
+    os.makedirs(args.save_path, exist_ok=True)
+
+    device = "cuda"
+    is_natural_image = args.dataset in ("gqa", "gqa_cot")
+    is_visualcot = args.dataset.endswith("_cot")
+
+    print(f"Loading {args.model_name} (dense, 64 layers, 64 heads, GQA 8:1)...")
+    model = Qwen3VLForConditionalGeneration.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+        device_map="auto",
+    ).eval()
+    print("Model loaded.")
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_path, trust_remote_code=True, padding_side="left", use_fast=True,
+    )
+
+    activations, hooks = register_hooks(model, is_qwen3vl=True)
+
+    layers = model.config.text_config.num_hidden_layers  # 64
+    heads = model.config.text_config.num_attention_heads  # 64
+    print(f"Config: {layers} layers, {heads} heads")
+
+    inspect_layers = None
+    if args.inspect_layers:
+        inspect_layers = [int(x) for x in args.inspect_layers.split(',')]
+        print(f"Will generate per-head maps for layers: {inspect_layers}")
+
+    images_path = args.images_path
+    if args.dataset == "textvqa" or args.dataset == "textvqa_cot":
+        images_path = images_path + "/textvqa/train_images"
+    elif args.dataset == "gqa_cot":
+        images_path = images_path + "/gqa/images"
+    elif args.dataset == "docvqa_cot":
+        images_path = images_path + "/DocVQA"
+
+    questions = [json.loads(q) for q in open(args.questions_path)]
+    # Visual-CoT jsonls are single-dataset → skip substring filter.
+    if not is_visualcot:
+        questions = [q for q in questions if args.dataset in os.path.join(images_path, q["image"])]
+    print(f"{args.dataset} questions: {len(questions)}")
+
+    sample = 0
+    for line in tqdm(questions):
+        activations.clear()
+
+        image_file = line["image"]
+        # Support LLaVA ("text"/"prompt") and Visual-CoT ("question") schemas.
+        qs = line.get("text") or line.get("question") or line.get("prompt") or ""
+
+        if is_natural_image:
+            # Always ensure the grounding-bbox instruction is present.
+            if 'Output grounding bounding box' not in qs:
+                if 'Answer the question using a single word or phrase.' in qs:
+                    qs = qs.replace('Answer the question using a single word or phrase.',
+                                    'Output grounding bounding box related to the question in JSON.')
+                else:
+                    qs = qs.rstrip() + ' Output grounding bounding box related to the question in JSON.'
+        else:
+            if 'Answer the question using a single word or phrase.' not in qs:
+                qs = qs.rstrip() + ' Answer the question using a single word or phrase.'
+
+        img_path = os.path.join(images_path, image_file)
+        if not os.path.exists(img_path):
+            continue
+
+        image = Image.open(img_path).convert('RGB')
+        processor.image_processor.max_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        processor.image_processor.min_pixels = 576 * PATCH_SIZE * PATCH_SIZE
+        image = expand2square(image, (127, 127, 127))
+
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": qs},
+                {"type": "image", "image": image},
+            ]
+        }]
+
+        inputs = processor.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True,
+            return_dict=True, return_tensors="pt"
+        )
+        inputs = inputs.to(device)
+
+        with torch.no_grad():
+            generated_ids = model.generate(**inputs, max_new_tokens=128, output_attentions=True)
+
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        answer = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        print(f"  [{sample}] Q: {qs[:60]}...")
+        print(f"       A: {answer[:60]}...")
+
+        # Build visual mask from <|vision_start|> ... <|vision_end|>
+        if activations.get('visual_mask', None) is None:
+            vision_start_id = processor.tokenizer.convert_tokens_to_ids('<|vision_start|>')
+            vision_end_id = processor.tokenizer.convert_tokens_to_ids('<|vision_end|>')
+            input_ids_list = inputs['input_ids'].tolist()[0]
+            pos = input_ids_list.index(vision_start_id) + 1
+            pos_end = input_ids_list.index(vision_end_id)
+            activations['visual_mask'] = torch.zeros(
+                (1, inputs['input_ids'].shape[1]), dtype=torch.bool, device=device)
+            activations['visual_mask'][0, pos:pos_end] = True
+
+        output_shape = inputs['image_grid_thw'].cpu().numpy().squeeze(0)[1:] / 2
+        output_shape = np.array(output_shape, dtype=np.int32)
+
+        attn_weights_map = get_attn_weights_map(activations, layers=layers, reduction='mean')
+
+        # Build a title_prefix from the question (truncated for readability on the figure).
+        qs_title = qs.replace("\n", " ").strip()
+        if len(qs_title) > 120:
+            qs_title = qs_title[:117] + "..."
+        title_prefix = f"[s{sample}] Q: {qs_title}\n"
+
+        # 1. Per-layer average attention map (all 64 layers in one grid)
+        visualize_token_wise_attention(
+            attn_weights_map,
+            qk_type='output2images',
+            attn_size=output_shape,
+            layers=layers,
+            rgb_image=image,
+            title_prefix=title_prefix,
+            save_path=f'{args.save_path}/{sample}_layer_attn.png'
+        )
+        print(f"       Saved: {sample}_layer_attn.png")
+
+        # 2. Per-head attention for specified layers (all 64 heads for each)
+        if inspect_layers:
+            for layer_idx in inspect_layers:
+                if layer_idx >= layers:
+                    print(f"       [WARN] layer {layer_idx} >= {layers}, skipping")
+                    continue
+                visualize_token_wise_attention(
+                    attn_weights_map,
+                    qk_type='output2images',
+                    avg_heads=False,
+                    rgb_image=image,
+                    layer_idx=layer_idx,
+                    attn_size=output_shape,
+                    layers=layers,
+                    title_prefix=title_prefix,
+                    save_path=f'{args.save_path}/{sample}_layer{layer_idx}_heads.png'
+                )
+                print(f"       Saved: {sample}_layer{layer_idx}_heads.png")
+
+        sample += 1
+        if sample >= args.debug_samples:
+            print('Debug samples limit reached.')
+            break
+
+    for h in hooks:
+        h.remove()
+    print(f"\nDone! {sample} samples saved to {args.save_path}/")
+    print(f"Next: inspect *_layer_attn.png for grounding layers in 64-layer stack,")
+    print(f"then re-run with --inspect_layers <L1,L2,...> for per-head maps,")
+    print(f"finally update required_heads/sink_heads in qwen3vl_debug_32b.py.")
+
+
+if __name__ == "__main__":
+    main()

QWENVL-PRIVATE/analysis/strip_image_root_inplace.py

@@ -0,0 +1,25 @@
+import json
+
+SRC = "/home/yuheng/datasets/llava_v1_5_mix665k_selected_qa.jsonl"
+DATA_ROOT = "/home/yuheng/datasets/"
+
+
+def main():
+    rows = []
+    with open(SRC, "r", encoding="utf-8") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            obj = json.loads(line)
+            img = obj.get("image", "")
+            if img.startswith(DATA_ROOT):
+                obj["image"] = img[len(DATA_ROOT) :]
+            rows.append(obj)
+
+    with open(SRC, "w", encoding="utf-8") as f:
+        for obj in rows:
+            f.write(json.dumps(obj, ensure_ascii=False) + "\n")
+
+
+if __name__ == "__main__":
+    main()

QWENVL-PRIVATE/analysis/tmp_plot.py

@@ -0,0 +1,84 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Benchmarks included: DocVQA, ChartQA, OCRBench, InfoVQA, TextVQA, V*, POPE
+models = ['Baseline-576', '+Ours (Full)', '+Ours-Gate', 'Baseline-1024']
+colors = ['#95a5a6', '#e74c3c', '#2ecc71', '#3498db']  # Grey, Red, Green, Blue
+markers = ['o', '^', '*', 's']
+
+# Data Structure: {Benchmark: (Score_List, Time_List)}
+# Order matches 'models' list
+data = {
+    'DocVQA': ([92.0, 93.6, 93.4, 93.9], [586, 806, 722, 799]),
+    'ChartQA': ([82.6, 85.6, 83.9, 84.4], [227, 396, 266, 355]),
+    'OCRBench': ([81.5, 84.1, 83.5, 83.1], [125, 193, 155, 172]),
+    'InfoVQA': ([70.2, 78.7, 77.2, 76.6], [251, 454, 355, 375]),
+    'TextVQA': ([81.1, 83.1, 82.7, 82.8], [436, 806, 506, 678]),
+    'V*': ([64.4, 76.4, 77.5, 70.2], [15, 29, 26, 24]),
+    'POPE': ([86.7, 88.6, 87.4, 86.7], [675, 1299, 788, 1115])
+}
+
+# Lists to store aggregated data
+avg_abs_scores = []
+avg_rel_times = []
+
+# Calculate Averages for each model
+num_models = len(models)
+for i in range(num_models):
+    total_score = 0
+    total_rel_time = 0
+    count = 0
+
+    for bench, (scores, times) in data.items():
+        # Sum absolute scores
+        total_score += scores[i]
+
+        # Sum relative time (Current Time / Baseline-576 Time)
+        base_time = times[0]
+        total_rel_time += times[i] / base_time
+
+        count += 1
+
+    avg_abs_scores.append(total_score / count)
+    avg_rel_times.append(total_rel_time / count)
+
+# Create Plot
+plt.figure(figsize=(11, 8))
+plt.grid(True, linestyle='--', alpha=0.6)
+
+# Plot Points
+for i, model in enumerate(models):
+    plt.scatter(avg_rel_times[i], avg_abs_scores[i],
+                color=colors[i], s=250, label=model, marker=markers[i], zorder=10, edgecolors='k')
+
+    # Dynamic annotation placement to avoid overlap
+    xytext_offset = (0, 15)
+    if 'Baseline-1024' in model:
+        xytext_offset = (0, -25)  # Move label down
+
+    plt.annotate(f"{model}\n({avg_rel_times[i]:.2f}x Time, {avg_abs_scores[i]:.1f}% Acc)",
+                 (avg_rel_times[i], avg_abs_scores[i]),
+                 xytext=xytext_offset, textcoords='offset points', ha='center',
+                 fontsize=10, fontweight='bold',
+                 bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.7, ec="none"))
+
+# Axis Limits (optional, for aesthetics)
+plt.xlim(0.9, 1.8)
+plt.ylim(78, 86)  # Adjust based on data range to zoom in on differences
+
+# Axis Labels
+plt.xlabel('Relative Time Cost (Normalized to Baseline-576)\n← Faster | Slower →', fontsize=12, fontweight='bold')
+plt.ylabel('Average Absolute Accuracy (%)\nHigher is Better ↑', fontsize=12, fontweight='bold')
+plt.title('Performance vs. Efficiency Trade-off\n(Average across 7 Benchmarks)', fontsize=14, pad=15)
+
+# Reference Lines
+# Vertical line at Baseline-576 time (1.0)
+plt.axvline(x=1.0, color='#95a5a6', linestyle=':', linewidth=2, label='Baseline Speed')
+
+# Legend
+plt.legend(loc='lower right', fontsize=11, frameon=True, framealpha=0.9, shadow=True)
+
+plt.tight_layout()
+plt.show()
+# plt.savefig('trade_off_plot_abs.png')
+# print("Plot saved as 'trade_off_plot_abs.png'")