Update Scholar Express with comprehensive Gradio app

- Add main app.py with DOLPHIN and Gemma 3n integration
- Include PDF processing, chat, and voice features
- Clean up demo files and deployment configs
- Update requirements for Hugging Face Spaces

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

gradio_add_voice.py → app.py

@@ -1,9 +1,3 @@
-"""
-DOLPHIN PDF Document AI - Local Gemma 3n Version
-Optimized for powerful GPU deployment with local models
-Features: AI-generated alt text for accessibility using local Gemma 3n
-"""
-
 import gradio as gr
 import json
 import markdown
@@ -51,14 +45,18 @@ try:
 except ImportError:
     pass
 
-# Warm up voice model if available
-if VOICE_DEPENDENCIES_AVAILABLE and voice_model:
+# Initialize voice model early to avoid NameError
+voice_model = None
+if VOICE_DEPENDENCIES_AVAILABLE:
     try:
+        print("Loading voice model...")
+        voice_model = Gemma3nInference(device='cuda' if torch.cuda.is_available() else 'cpu')
         print("Warming up voice model...")
         voice_model.warm_up()
-        print("✅ Voice model warmed up successfully")
+        print("✅ Voice model loaded and warmed up successfully")
     except Exception as e:
-        print(f"⚠️ Voice model warm-up failed: {e}")
+        print(f"⚠️ Voice model initialization failed: {e}")
+        voice_model = None
 
 
 class DOLPHIN:
@@ -520,16 +518,7 @@ OUT_RATE = 24000
 OUT_SAMPLE_WIDTH = 2
 OUT_CHUNK = 20 * 4096
 
-# Initialize voice inference model if available
-voice_model = None
-if VOICE_DEPENDENCIES_AVAILABLE:
-    try:
-        print("Loading voice model for Talk with Gemma...")
-        voice_model = Gemma3nInference(device='cuda' if torch.cuda.is_available() else 'cpu')
-        print("✅ Voice model loaded successfully")
-    except Exception as e:
-        print(f"❌ Error loading voice model: {e}")
-        VOICE_DEPENDENCIES_AVAILABLE = False
+# Voice model already initialized earlier in the file
 
 @dataclass
 class VoiceAppState:
@@ -643,8 +632,8 @@ def generate_voice_response(state: VoiceAppState):
                 audio_array = audio_array.reshape((-1, 2))
             
             # Update conversation history
-            state.conversation.append({"role": "user", "content": {"path": temp_audio_path, "mime_type": "audio/wav"}})
-            state.conversation.append({"role": "assistant", "content": {"text": text_response}})
+            state.conversation.append({"role": "user", "content": f"[Audio message]"})
+            state.conversation.append({"role": "assistant", "content": text_response})
             
             return (audio_segment.frame_rate, audio_array), VoiceAppState(conversation=state.conversation)
             
@@ -697,7 +686,7 @@ def create_embeddings(chunks):
 def retrieve_relevant_chunks(question, chunks, embeddings, top_k=3):
     """Retrieve most relevant chunks for a question"""
     if embedding_model is None or embeddings is None:
-        return chunks[:3]  # Fallback to first 3 chunks
+        return chunks[:3] =
     
     try:
         question_embedding = embedding_model.encode([question], show_progress_bar=False)
@@ -892,6 +881,7 @@ with gr.Blocks(
             chatbot = gr.Chatbot(
                 value=[],
                 height=500,
+                type='messages',
                 elem_classes="chatbot-container",
                 placeholder="Your conversation will appear here once you process a document..."
             )

chat.py

@@ -1,8 +1,3 @@
-""" 
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
 import os
 import warnings
 from collections import OrderedDict

demo/element_imgs/markdown/table_1.md

@@ -1,2 +0,0 @@
-<table><tr><td></td><td></td><td>100-class (top-1 acc.)</td><td>1000-class (top-1 acc.)</td></tr><tr><td colspan="2">4096-d (float)</td><td>77.1 ± 1.5</td><td>65.0</td></tr><tr><td rowspan="3">1024 bits</td><td>BP</td><td>72.9 ± 1.3</td><td>58.1</td></tr><tr><td>CBE</td><td>73.0 ± 1.3</td><td>59.2</td></tr><tr><td>SP</td><td>73.8 ± 1.3</td><td>60.1</td></tr><tr><td rowspan="4">4096 bits</td><td>threshold [1]</td><td>73.5 ± 1.4</td><td>59.1</td></tr><tr><td>BP</td><td>76.0 ± 1.5</td><td>63.2</td></tr><tr><td>CBE</td><td>75.9 ± 1.4</td><td>63.0</td></tr><tr><td>SP</td><td>76.3 ± 1.5</td><td>63.3</td></tr><tr><td>8192 bits</td><td>SP</td><td>76.8 ± 1.4</td><td>64.2</td></tr><tr><td>16384 bits</td><td>SP</td><td>77.1 ± 1.6</td><td>64.5</td></tr></table>
-

demo/element_imgs/recognition_json/table_1.json

@@ -1,6 +0,0 @@
-[
-  {
-    "label": "tab",
-    "text": "<table><tr><td></td><td></td><td>100-class (top-1 acc.)</td><td>1000-class (top-1 acc.)</td></tr><tr><td colspan=\"2\">4096-d (float)</td><td>77.1 ± 1.5</td><td>65.0</td></tr><tr><td rowspan=\"3\">1024 bits</td><td>BP</td><td>72.9 ± 1.3</td><td>58.1</td></tr><tr><td>CBE</td><td>73.0 ± 1.3</td><td>59.2</td></tr><tr><td>SP</td><td>73.8 ± 1.3</td><td>60.1</td></tr><tr><td rowspan=\"4\">4096 bits</td><td>threshold [1]</td><td>73.5 ± 1.4</td><td>59.1</td></tr><tr><td>BP</td><td>76.0 ± 1.5</td><td>63.2</td></tr><tr><td>CBE</td><td>75.9 ± 1.4</td><td>63.0</td></tr><tr><td>SP</td><td>76.3 ± 1.5</td><td>63.3</td></tr><tr><td>8192 bits</td><td>SP</td><td>76.8 ± 1.4</td><td>64.2</td></tr><tr><td>16384 bits</td><td>SP</td><td>77.1 ± 1.6</td><td>64.5</td></tr></table>"
-  }
-]

demo/page_imgs/markdown/test_page3.md

@@ -1,22 +0,0 @@
-![Figure](figures/test_page3_figure_000.png)
-
-Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several attention layers running in parallel.
-
-query with all keys, divide each by $\sqrt{d_k}$ , and apply a softmax function to obtain the weights on the values.
-
-In practice, we compute the attention function on a set of queries simultaneously, packed together into a matrix $Q$ . The keys and values are also packed together into matrices $K$ and $V$ . We compute the matrix of outputs as: $$ \\     \text{Attention}(Q, K, V) = \mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V \\ $$
-
-The two most commonly used attention functions are additive attention [2] , and dot-product (multiplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of $\frac{1}{\sqrt{d_k}}$ . Additive attention computes the compatibility function using a feed-forward network with a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.
-
-While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of $d_k$ [ 3 ] . We suspect that for large values of $d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients 4 To counteract this effect, we scale the dot products by $\frac{1}{\sqrt{d_k}}$ .
-
-3.2.2 Multi-Head Attention
-
-Instead of performing a single attention function with $d_{\text{model}}$ -dimensional keys, values and queries, we found it beneficial to linearly project the queries, keys and values $h$ times with different, learned linear projections to $d_k$ , $d_k$ and $d_v$ dimensions, respectively. On each of these projected versions of queries, keys and values we then perform the attention function in parallel, yielding $d_v$ -dimensional output values. These are concatenated and once again projected, resulting in the final values, as depicted in Figure 2 .
-
-Multi­head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.
-
-${ }^{4}$ To illustrate why the dot products get large, assume that the components of $q$ and $k$ are independent random variables with mean 0 and variance 1 . Then their dot product, $q \cdot k=\sum_{i=1}^{d_{k}} q_{i} k_{i}$, has mean 0 and variance $d_{k}$.
-
-4
-

demo/page_imgs/recognition_json/page_1.json

@@ -1,178 +0,0 @@
-[
-  {
-    "label": "title",
-    "bbox": [
-      271,
-      188,
-      1194,
-      221
-    ],
-    "text": "LLaMA: Open and Efficient Foundation Language Models",
-    "reading_order": 0
-  },
-  {
-    "label": "author",
-    "bbox": [
-      313,
-      272,
-      1154,
-      317
-    ],
-    "text": "Hugo Touvron; Thibaut Lavril*, Gautier Izacard*, Xavier Martinet",
-    "reading_order": 1
-  },
-  {
-    "label": "para",
-    "bbox": [
-      269,
-      317,
-      1201,
-      425
-    ],
-    "text": "Marie-Anne Lachaux, Timothee Lacroix, Baptiste Rozière, Naman Goyal\nEric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin\nEdouard Grave*Guillaume Lample*",
-    "reading_order": 2
-  },
-  {
-    "label": "para",
-    "bbox": [
-      685,
-      440,
-      795,
-      482
-    ],
-    "text": "Meta AI",
-    "reading_order": 3
-  },
-  {
-    "label": "sec",
-    "bbox": [
-      376,
-      524,
-      502,
-      565
-    ],
-    "text": "\\begin{abstract}",
-    "reading_order": 4
-  },
-  {
-    "label": "para",
-    "bbox": [
-      209,
-      586,
-      675,
-      946
-    ],
-    "text": "We introduce LLaMA, a collection of founda-\ntion language models ranging from 7B to 65B\nparameters. We train our models on trillions\nof tokens, and show that it is possible to train\nstate-of-the-art models using publicly avail-\nable datasets exclusively, without resorting\nto proprietary and inaccessible datasets. In\nparticular, LLaMA-13B outperforms GPT-3\n(175B) on most benchmarks, and LLaMA-\n65B is competitive with the best models,\nChinchilla-70B and PaLM-540B. We release\nall our models to the research community $^1$ .",
-    "reading_order": 5
-  },
-  {
-    "label": "sec",
-    "bbox": [
-      167,
-      964,
-      376,
-      1006
-    ],
-    "text": "1 Introduction",
-    "reading_order": 6
-  },
-  {
-    "label": "para",
-    "bbox": [
-      167,
-      1027,
-      718,
-      1498
-    ],
-    "text": "Large Languages Models (LLMs) trained on mas-\nsive corpora of texts have shown their ability to per-\nform new tasks from textual instructions or from a\nfew examples ( Brown et al. , 2020 ) . These few-shot\nproperties first appeared when scaling models to a\nsufficient size ( Kaplan et al. , 2020 ) , resulting in a\nline of work that focuses on further scaling these\nmodels ( Chowdhery et al. , 2022 ; Rae et al. , 2021 ) .\nThese efforts are based on the assumption that\nmore parameters will lead to better performance.\nHowever, recent work from Hoffmann et al. ( 2022 )\nshows that, for a given compute budget, the best\nperformances are not achieved by the largest mod-\nels, but by smaller models trained on more data.",
-    "reading_order": 7
-  },
-  {
-    "label": "para",
-    "bbox": [
-      167,
-      1506,
-      717,
-      1844
-    ],
-    "text": "The objective of the scaling laws from Hoff-\nmann et al. ( 2022 ) is to determine how to best\nscale the dataset and model sizes for a particular\ntraining compute budget. However, this objective\ndisregards the inference budget, which becomes\ncritical when serving a language model at scale.\nIn this context, given a target level of performance,\nthe preferred model is not the fastest to train but the\nfastest at inference, and although it may be cheaper\nto train a large model to reach a certain level of",
-    "reading_order": 8
-  },
-  {
-    "label": "para",
-    "bbox": [
-      753,
-      539,
-      1304,
-      734
-    ],
-    "text": "performance, a smaller one trained longer will\nultimately be cheaper at inference. For instance,\nalthough Hoffmann et al. ( 2022 ) recommends\ntraining a 10B model on 200B tokens, we find\nthat the performance of a 7B model continues to\nimprove even after 1T tokens.",
-    "reading_order": 9
-  },
-  {
-    "label": "para",
-    "bbox": [
-      753,
-      769,
-      1305,
-      1236
-    ],
-    "text": "The focus of this work is to train a series of\nlanguage models that achieve the best possible per-\nformance at various inference budgets, by training\non more tokens than what is typically used. The\nresulting models, called LLaMA , ranges from 7B\nto 65B parameters with competitive performance\ncompared to the best existing LLMs. For instance,\nLLaMA-13B outperforms GPT-3 on most bench-\nmarks, despite being 10 $\\times$ smaller. We believe that\nthis model will help democratize the access and\nstudy of LLMs, since it can be run on a single GPU.\nAt the higher-end of the scale, our 65B-parameter\nmodel is also competitive with the best large lan-\nguage models such as Chinchilla or PaLM-540B.",
-    "reading_order": 10
-  },
-  {
-    "label": "para",
-    "bbox": [
-      753,
-      1257,
-      1305,
-      1601
-    ],
-    "text": "Unlike Chinchilla, PaLM, or GPT-3, we only\nuse publicly available data, making our work com-\npatible with open-sourcing, while most existing\nmodels rely on data which is either not publicly\navailable or undocumented (e.g. “ Books – 2TB ” or\n“ Social media conversations ” ). There exist some\nexceptions, notably OPT ( Zhang et al. , 2022 ) ,\nGPT-NeoX ( Black et al. , 2022 ) , BLOOM ( Scao\net al. , 2022 ) and GLM ( Zeng et al. , 2022 ) , but none\nthat are competitive with PaLM-62B or Chinchilla.",
-    "reading_order": 11
-  },
-  {
-    "label": "para",
-    "bbox": [
-      753,
-      1634,
-      1304,
-      1933
-    ],
-    "text": "In the rest of this paper, we present an overview\nof the modifications we made to the transformer\narchitecture ( Vaswani et al. , 2017 ) , as well as our\ntraining method. We then report the performance of\nour models and compare with others LLMs on a set\nof standard benchmarks. Finally, we expose some\nof the biases and toxicity encoded in our models,\nusing some of the most recent benchmarks from\nthe responsible AI community.",
-    "reading_order": 12
-  },
-  {
-    "label": "fnote",
-    "bbox": [
-      167,
-      1844,
-      712,
-      1907
-    ],
-    "text": "* Equal contribution.\nCorrespondence:\n{htouvron\nthibautlav,gizacard,egrave,glample}@meta.com",
-    "reading_order": 13
-  },
-  {
-    "label": "fnote",
-    "bbox": [
-      209,
-      1907,
-      632,
-      1931
-    ],
-    "text": "https://github.com/facebookresearch/llama",
-    "reading_order": 14
-  },
-  {
-    "label": "watermark",
-    "bbox": [
-      20,
-      649,
-      83,
-      1530
-    ],
-    "text": "arXiv:2302.1397lvl [cs.CL] 27 Feb 2023",
-    "reading_order": 15
-  }
-]

demo/page_imgs/recognition_json/test_page.json

@@ -1,47 +0,0 @@
-[
-  {
-    "label": "header",
-    "bbox": [
-      291,
-      90,
-      675,
-      120
-    ],
-    "text": "Scaled Dot-Product Attention",
-    "reading_order": 0
-  },
-  {
-    "label": "fig",
-    "text": "![Figure](figures/test_page_figure_001.png)",
-    "figure_path": "figures/test_page_figure_001.png",
-    "bbox": [
-      1274,
-      105,
-      1536,
-      627
-    ],
-    "reading_order": 1
-  },
-  {
-    "label": "cap",
-    "bbox": [
-      168,
-      719,
-      1413,
-      789
-    ],
-    "text": "Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several\nattention layers running in parallel.",
-    "reading_order": 2
-  },
-  {
-    "label": "para",
-    "bbox": [
-      168,
-      858,
-      1413,
-      934
-    ],
-    "text": "query with all keys, divide each by $\\sqrt{d_{k}}$, and apply a softmax function to obtain the weights on the\nvalues.",
-    "reading_order": 3
-  }
-]

demo/page_imgs/recognition_json/test_page2.json

@@ -1,102 +0,0 @@
-[
-  {
-    "label": "fig",
-    "text": "![Figure](figures/test_page2_figure_000.png)",
-    "figure_path": "figures/test_page2_figure_000.png",
-    "bbox": [
-      394,
-      117,
-      897,
-      837
-    ],
-    "reading_order": 0
-  },
-  {
-    "label": "cap",
-    "bbox": [
-      445,
-      852,
-      856,
-      873
-    ],
-    "text": "Figure 1: The Transformer - model architecture",
-    "reading_order": 1
-  },
-  {
-    "label": "para",
-    "bbox": [
-      218,
-      920,
-      1086,
-      1044
-    ],
-    "text": "wise fully connected feed-forward network. We employ a residual connection [ 10 ] around each of\nthe two sub-layers, followed by layer normalization [ 1 ] . That is, the output of each sub-layer is\n$\\mathrm{LayerNorm}(x+\\mathrm{Sublayer}(x))$ , where $\\mathrm{Sublayer}(x)$ is the function implemented by the sub-layer\nitself. To facilitate these residual connections, all sub-layers in the model, as well as the embedding\nlayers, produce outputs of dimension $d_{\\text{model}}=512$ .",
-    "reading_order": 2
-  },
-  {
-    "label": "para",
-    "bbox": [
-      218,
-      1071,
-      1085,
-      1244
-    ],
-    "text": "The The decoder is also composed of a stack of $N=6$ identical layers. In addition to the two\nsub-layers in each encoder layer, the decoder inserts a third sub-layer, which performs multi-head\nattention over the output of the encoder stack. Similar to the encoder, we employ residual connections\naround each of the sub-layers, followed by layer normalization. We also modify the self-attention\nsub-layer in the decoder stack to prevent positions from attending to subsequent positions. This\nmasking, combined with fact that the output embeddings are offset by one position, ensures that the\npredictions for position $i$ can depend only on the known outputs at positions less than $i$ .",
-    "reading_order": 3
-  },
-  {
-    "label": "sub_sec",
-    "bbox": [
-      226,
-      1283,
-      344,
-      1305
-    ],
-    "text": "3.2 Attention",
-    "reading_order": 4
-  },
-  {
-    "label": "para",
-    "bbox": [
-      218,
-      1322,
-      1087,
-      1422
-    ],
-    "text": "An attention function can be described as mapping a query and a set of key-value pairs to an output,\nwhere the query, keys, values, and output are all vectors. The output is computed as a weighted sum\nof the values, where the weight assigned to each value is computed by a compatibility function of the\nquery with the corresponding key.",
-    "reading_order": 5
-  },
-  {
-    "label": "sub_sub_sec",
-    "bbox": [
-      218,
-      1456,
-      562,
-      1474
-    ],
-    "text": "3.2.1 Scaled Dot-Product Attention",
-    "reading_order": 6
-  },
-  {
-    "label": "para",
-    "bbox": [
-      218,
-      1498,
-      1085,
-      1546
-    ],
-    "text": "We call our particular attention \"Scaled Dot-Product Attention\" (Figure 2 ). The input consists of\nqueries and keys of dimension $d_k$ , and values of dimension $d_v$ . We compute the dot products of the",
-    "reading_order": 7
-  },
-  {
-    "label": "foot",
-    "bbox": [
-      646,
-      1590,
-      662,
-      1607
-    ],
-    "text": "3",
-    "reading_order": 8
-  }
-]

demo/page_imgs/recognition_json/test_page3.json

@@ -1,124 +0,0 @@
-[
-  {
-    "label": "fig",
-    "text": "![Figure](figures/test_page3_figure_000.png)",
-    "figure_path": "figures/test_page3_figure_000.png",
-    "bbox": [
-      331,
-      134,
-      984,
-      489
-    ],
-    "reading_order": 0
-  },
-  {
-    "label": "cap",
-    "bbox": [
-      198,
-      554,
-      1065,
-      603
-    ],
-    "text": "Figure 2: (left) Scaled Dot-Product Attention. (right) Multi-Head Attention consists of several\nattention layers running in parallel.",
-    "reading_order": 1
-  },
-  {
-    "label": "para",
-    "bbox": [
-      198,
-      652,
-      1065,
-      701
-    ],
-    "text": "query with all keys, divide each by $\\sqrt{d_k}$ , and apply a softmax function to obtain the weights on the\nvalues.",
-    "reading_order": 2
-  },
-  {
-    "label": "para",
-    "bbox": [
-      198,
-      715,
-      1065,
-      881
-    ],
-    "text": "In practice, we compute the attention function on a set of queries simultaneously, packed together\ninto a matrix $Q$ . The keys and values are also packed together into matrices $K$ and $V$ . We compute\nthe matrix of outputs as:\n\\[\n    \\text{Attention}(Q, K, V) = \\mathrm{softmax}(\\frac{QK^T}{\\sqrt{d_k}})V\n\\]",
-    "reading_order": 3
-  },
-  {
-    "label": "para",
-    "bbox": [
-      198,
-      913,
-      1068,
-      1060
-    ],
-    "text": "The two most commonly used attention functions are additive attention [2] , and dot-product (multi-\nplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor\nof $\\frac{1}{\\sqrt{d_k}}$ . Additive attention computes the compatibility function using a feed-forward network with\na single hidden layer. While the two are similar in theoretical complexity, dot-product attention is\nmuch faster and more space-efficient in practice, since it can be implemented using highly optimized\nmatrix multiplication code.",
-    "reading_order": 4
-  },
-  {
-    "label": "para",
-    "bbox": [
-      198,
-      1074,
-      1066,
-      1175
-    ],
-    "text": "While for small values of $d_k$ the two mechanisms perform similarly, additive attention outperforms\ndot product attention without scaling for larger values of $d_k$ [ 3 ] . We suspect that for large values of\n$d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions where it has\nextremely small gradients 4 To counteract this effect, we scale the dot products by $\\frac{1}{\\sqrt{d_k}}$ .",
-    "reading_order": 5
-  },
-  {
-    "label": "sub_sub_sec",
-    "bbox": [
-      198,
-      1207,
-      467,
-      1225
-    ],
-    "text": "3.2.2 Multi-Head Attention",
-    "reading_order": 6
-  },
-  {
-    "label": "para",
-    "bbox": [
-      198,
-      1253,
-      1067,
-      1395
-    ],
-    "text": "Instead of performing a single attention function with $d_{\\text{model}}$ -dimensional keys, values and queries,\nwe found it beneficial to linearly project the queries, keys and values $h$ times with different, learned\nlinear projections to $d_k$ , $d_k$ and $d_v$ dimensions, respectively. On each of these projected versions of\nqueries, keys and values we then perform the attention function in parallel, yielding $d_v$ -dimensional\noutput values. These are concatenated and once again projected, resulting in the final values, as\ndepicted in Figure 2 .",
-    "reading_order": 7
-  },
-  {
-    "label": "para",
-    "bbox": [
-      198,
-      1403,
-      1065,
-      1453
-    ],
-    "text": "Multi­head attention allows the model to jointly attend to information from different representation\nsubspaces at different positions. With a single attention head, averaging inhibits this.",
-    "reading_order": 8
-  },
-  {
-    "label": "fnote",
-    "bbox": [
-      198,
-      1485,
-      1065,
-      1535
-    ],
-    "text": "${ }^{4}$ To illustrate why the dot products get large, assume that the components of $q$ and $k$ are independent random\nvariables with mean 0 and variance 1 . Then their dot product, $q \\cdot k=\\sum_{i=1}^{d_{k}} q_{i} k_{i}$, has mean 0 and variance $d_{k}$.",
-    "reading_order": 9
-  },
-  {
-    "label": "foot",
-    "bbox": [
-      625,
-      1578,
-      641,
-      1599
-    ],
-    "text": "4",
-    "reading_order": 10
-  }
-]

demo_element.py

@@ -1,129 +0,0 @@
-""" 
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
-import argparse
-import glob
-import os
-
-from omegaconf import OmegaConf
-from PIL import Image
-
-from chat import DOLPHIN
-from utils.utils import *
-
-
-def process_element(image_path, model, element_type, save_dir=None):
-    """Process a single element image (text, table, formula)
-
-    Args:
-        image_path: Path to the element image
-        model: DOLPHIN model instance
-        element_type: Type of element ('text', 'table', 'formula')
-        save_dir: Directory to save results (default: same as input directory)
-
-    Returns:
-        Parsed content of the element and recognition results
-    """
-    # Load and prepare image
-    pil_image = Image.open(image_path).convert("RGB")
-    pil_image = crop_margin(pil_image)
-
-    # Select appropriate prompt based on element type
-    if element_type == "table":
-        prompt = "Parse the table in the image."
-        label = "tab"
-    elif element_type == "formula":
-        prompt = "Read text in the image."
-        label = "formula"
-    else:  # Default to text
-        prompt = "Read text in the image."
-        label = "text"
-
-    # Process the element
-    result = model.chat(prompt, pil_image)
-
-    # Create recognition result in the same format as the document parser
-    recognition_result = [
-        {
-            "label": label,
-            "text": result.strip(),
-        }
-    ]
-
-    # Save results if save_dir is provided
-    if save_dir:
-        save_outputs(recognition_result, image_path, save_dir)
-        print(f"Results saved to {save_dir}")
-
-    return result, recognition_result
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Element-level processing using DOLPHIN model")
-    parser.add_argument("--config", default="./config/Dolphin.yaml", help="Path to configuration file")
-    parser.add_argument("--input_path", type=str, required=True, help="Path to input image or directory of images")
-    parser.add_argument(
-        "--element_type",
-        type=str,
-        choices=["text", "table", "formula"],
-        default="text",
-        help="Type of element to process (text, table, formula)",
-    )
-    parser.add_argument(
-        "--save_dir",
-        type=str,
-        default=None,
-        help="Directory to save parsing results (default: same as input directory)",
-    )
-    parser.add_argument("--print_results", action="store_true", help="Print recognition results to console")
-    args = parser.parse_args()
-
-    # Load Model
-    config = OmegaConf.load(args.config)
-    model = DOLPHIN(config)
-
-    # Set save directory
-    save_dir = args.save_dir or (
-        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
-    )
-    setup_output_dirs(save_dir)
-
-    # Collect Images
-    if os.path.isdir(args.input_path):
-        image_files = []
-        for ext in [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"]:
-            image_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
-        image_files = sorted(image_files)
-    else:
-        if not os.path.exists(args.input_path):
-            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
-        image_files = [args.input_path]
-
-    total_samples = len(image_files)
-    print(f"\nTotal samples to process: {total_samples}")
-
-    # Process images one by one
-    for image_path in image_files:
-        print(f"\nProcessing {image_path}")
-        try:
-            result, recognition_result = process_element(
-                image_path=image_path,
-                model=model,
-                element_type=args.element_type,
-                save_dir=save_dir,
-            )
-
-            if args.print_results:
-                print("\nRecognition result:")
-                print(result)
-                print("-" * 40)
-
-        except Exception as e:
-            print(f"Error processing {image_path}: {str(e)}")
-            continue
-
-
-if __name__ == "__main__":
-    main()

demo_element_hf.py

@@ -1,8 +1,3 @@
-"""
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
 import argparse
 import glob
 import os

demo_page.py

@@ -1,247 +0,0 @@
-""" 
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
-import argparse
-import glob
-import os
-
-import cv2
-from omegaconf import OmegaConf
-from PIL import Image
-
-from chat import DOLPHIN
-from utils.utils import *
-
-
-def process_document(document_path, model, save_dir, max_batch_size):
-    """Parse documents - Handles both images and PDFs"""
-    file_ext = os.path.splitext(document_path)[1].lower()
-    
-    if file_ext == '.pdf':
-        # Process PDF file
-        # Convert PDF to images
-        images = convert_pdf_to_images(document_path)
-        if not images:
-            raise Exception(f"Failed to convert PDF {document_path} to images")
-        
-        all_results = []
-        
-        # Process each page
-        for page_idx, pil_image in enumerate(images):
-            print(f"Processing page {page_idx + 1}/{len(images)}")
-            
-            # Generate output name for this page
-            base_name = os.path.splitext(os.path.basename(document_path))[0]
-            page_name = f"{base_name}_page_{page_idx + 1:03d}"
-            
-            # Process this page (don't save individual page results)
-            json_path, recognition_results = process_single_image(
-                pil_image, model, save_dir, page_name, max_batch_size, save_individual=False
-            )
-            
-            # Add page information to results
-            page_results = {
-                "page_number": page_idx + 1,
-                "elements": recognition_results
-            }
-            all_results.append(page_results)
-        
-        # Save combined results for multi-page PDF
-        combined_json_path = save_combined_pdf_results(all_results, document_path, save_dir)
-        
-        return combined_json_path, all_results
-
-    else:
-        # Process regular image file
-        pil_image = Image.open(document_path).convert("RGB")
-        base_name = os.path.splitext(os.path.basename(document_path))[0]
-        return process_single_image(pil_image, model, save_dir, base_name, max_batch_size)
-
-
-def process_single_image(image, model, save_dir, image_name, max_batch_size, save_individual=True):
-    """Process a single image (either from file or converted from PDF page)
-    
-    Args:
-        image: PIL Image object
-        model: DOLPHIN model instance
-        save_dir: Directory to save results
-        image_name: Name for the output file
-        max_batch_size: Maximum batch size for processing
-        save_individual: Whether to save individual results (False for PDF pages)
-        
-    Returns:
-        Tuple of (json_path, recognition_results)
-    """
-    # Stage 1: Page-level layout and reading order parsing
-    layout_output = model.chat("Parse the reading order of this document.", image)
-
-    # Stage 2: Element-level content parsing
-    padded_image, dims = prepare_image(image)
-    recognition_results = process_elements(layout_output, padded_image, dims, model, max_batch_size, save_dir, image_name)
-
-    # Save outputs only if requested (skip for PDF pages)
-    json_path = None
-    if save_individual:
-        # Create a dummy image path for save_outputs function
-        dummy_image_path = f"{image_name}.jpg"  # Extension doesn't matter, only basename is used
-        json_path = save_outputs(recognition_results, dummy_image_path, save_dir)
-
-    return json_path, recognition_results
-
-
-def process_elements(layout_results, padded_image, dims, model, max_batch_size, save_dir=None, image_name=None):
-    """Parse all document elements with parallel decoding"""
-    layout_results = parse_layout_string(layout_results)
-
-    text_table_elements = []  # Elements that need processing
-    figure_results = []  # Figure elements (no processing needed)
-    previous_box = None
-    reading_order = 0
-
-    # Collect elements for processing
-    for bbox, label in layout_results:
-        try:
-            # Adjust coordinates
-            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
-                bbox, padded_image, dims, previous_box
-            )
-
-            # Crop and parse element
-            cropped = padded_image[y1:y2, x1:x2]
-            if cropped.size > 0 and cropped.shape[0] > 3 and cropped.shape[1] > 3:
-                if label == "fig":
-                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-                    
-                    figure_filename = save_figure_to_local(pil_crop, save_dir, image_name, reading_order)
-                    
-                    # For figure regions, store relative path instead of base64
-                    figure_results.append(
-                        {
-                            "label": label,
-                            "text": f"![Figure](figures/{figure_filename})",
-                            "figure_path": f"figures/{figure_filename}",
-                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                            "reading_order": reading_order,
-                        }
-                    )
-                else:
-                    # For text or table regions, prepare for parsing
-                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-                    prompt = "Parse the table in the image." if label == "tab" else "Read text in the image."
-                    text_table_elements.append(
-                        {
-                            "crop": pil_crop,
-                            "prompt": prompt,
-                            "label": label,
-                            "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                            "reading_order": reading_order,
-                        }
-                    )
-
-            reading_order += 1
-
-        except Exception as e:
-            print(f"Error processing bbox with label {label}: {str(e)}")
-            continue
-
-    # Parse text/table elements in parallel
-    recognition_results = figure_results
-    if text_table_elements:
-        crops_list = [elem["crop"] for elem in text_table_elements]
-        prompts_list = [elem["prompt"] for elem in text_table_elements]
-
-        # Inference in batch
-        batch_results = model.chat(prompts_list, crops_list, max_batch_size=max_batch_size)
-
-        # Add batch results to recognition_results
-        for i, result in enumerate(batch_results):
-            elem = text_table_elements[i]
-            recognition_results.append(
-                {
-                    "label": elem["label"],
-                    "bbox": elem["bbox"],
-                    "text": result.strip(),
-                    "reading_order": elem["reading_order"],
-                }
-            )
-
-    # Sort elements by reading order
-    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
-
-    return recognition_results
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Document parsing based on DOLPHIN")
-    parser.add_argument("--config", default="./config/Dolphin.yaml", help="Path to configuration file")
-    parser.add_argument("--input_path", type=str, default="./demo", help="Path to input image/PDF or directory of files")
-    parser.add_argument(
-        "--save_dir",
-        type=str,
-        default=None,
-        help="Directory to save parsing results (default: same as input directory)",
-    )
-    parser.add_argument(
-        "--max_batch_size",
-        type=int,
-        default=4,
-        help="Maximum number of document elements to parse in a single batch (default: 4)",
-    )
-    args = parser.parse_args()
-
-    # Load Model
-    config = OmegaConf.load(args.config)
-    model = DOLPHIN(config)
-
-    # Collect Document Files (images and PDFs)
-    if os.path.isdir(args.input_path):
-        # Support both image and PDF files
-        file_extensions = [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG", ".pdf", ".PDF"]
-        
-        document_files = []
-        for ext in file_extensions:
-            document_files.extend(glob.glob(os.path.join(args.input_path, f"*{ext}")))
-        document_files = sorted(document_files)
-    else:
-        if not os.path.exists(args.input_path):
-            raise FileNotFoundError(f"Input path {args.input_path} does not exist")
-        
-        # Check if it's a supported file type
-        file_ext = os.path.splitext(args.input_path)[1].lower()
-        supported_exts = ['.jpg', '.jpeg', '.png', '.pdf']
-        
-        if file_ext not in supported_exts:
-            raise ValueError(f"Unsupported file type: {file_ext}. Supported types: {supported_exts}")
-        
-        document_files = [args.input_path]
-
-    save_dir = args.save_dir or (
-        args.input_path if os.path.isdir(args.input_path) else os.path.dirname(args.input_path)
-    )
-    setup_output_dirs(save_dir)
-
-    total_samples = len(document_files)
-    print(f"\nTotal files to process: {total_samples}")
-
-    # Process All Document Files
-    for file_path in document_files:
-        print(f"\nProcessing {file_path}")
-        try:
-            json_path, recognition_results = process_document(
-                document_path=file_path,
-                model=model,
-                save_dir=save_dir,
-                max_batch_size=args.max_batch_size,
-            )
-
-            print(f"Processing completed. Results saved to {save_dir}")
-
-        except Exception as e:
-            print(f"Error processing {file_path}: {str(e)}")
-            continue
-
-
-if __name__ == "__main__":
-    main()

demo_page_hf.py

@@ -1,8 +1,3 @@
-""" 
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
 import argparse
 import glob
 import os

deployment/ReadMe.md

@@ -1,12 +0,0 @@
-<h1 align="center">
-🚀 Dolphin Inference/Serving
-</h1>
-
-## vLLM
-> [Doc](./vllm/ReadMe.md)
-
-## TensorRT-LLM
-> [Doc](./tensorrt_llm/ReadMe.md)
-
-## Others
-

deployment/tensorrt_llm/ReadMe.md

@@ -1,89 +0,0 @@
-<h1 align="center">
-🚀 Dolphin TensorRT-LLM Demo
-</h1>
-
-## ✅ Introduction
-The Dolphin model employs a **Swin Encoder + MBart Decoder** architecture. In the HuggingFace Transformers [Config](https://huggingface.co/ByteDance/Dolphin/blob/main/config.json), 
-its architectures field is specified as "VisionEncoderDecoderModel". **Dolphin**, **[Nougat](https://huggingface.co/docs/transformers/model_doc/nougat)**, and **[Donut](https://huggingface.co/docs/transformers/model_doc/donut)** share the same model architecture. TensorRT-LLM has already supported the Nougat model. 
-Following Nougat's conversion script, we have successfully implemented Dolphin on TensorRT-LLM. 
-
-**Note:** [prompt_ids](./dolphin_runner.py#L120) MUST be of **int32** type, otherwise TensorRT-LLM will produce incorrect results.
-
-## 🛠️ Installation
-> We only test TensorRT-LLM 0.18.1 on Linux.
-
-https://nvidia.github.io/TensorRT-LLM/0.18.1/installation/linux.html
-
-
-## ⚡ Offline Inference
-```
-export MODEL_NAME="Dolphin"
-
-# predict elements reading order
-python run_dolphin.py \
-    --batch_size 1 \
-    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
-    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
-    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
-    --max_new_tokens 4096 \
-    --repetition_penalty 1.0 \
-    --input_text "Parse the reading order of this document." \
-    --image_path "../../demo/page_imgs/page_1.jpeg"
-
-# recognize text/latex
-python run_dolphin.py \
-    --batch_size 1 \
-    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
-    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
-    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
-    --max_new_tokens 4096 \
-    --repetition_penalty 1.0 \
-    --input_text "Read text in the image." \
-    --image_path "../../demo/element_imgs/block_formula.jpeg"
-
-
-python run_dolphin.py \
-    --batch_size 1 \
-    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
-    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
-    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
-    --max_new_tokens 4096 \
-    --repetition_penalty 1.0 \
-    --input_text "Read text in the image." \
-    --image_path "../../demo/element_imgs/para_1.jpg"
-
-# recognize table
-python run_dolphin.py \
-    --batch_size 1 \
-    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
-    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
-    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
-    --max_new_tokens 4096 \
-    --repetition_penalty 1.0 \
-    --input_text "Parse the table in the image." \
-    --image_path "../../demo/element_imgs/table_1.jpeg"
-```
-
-
-## ⚡ Online Inference
-```
-# 1. Start Api Server
-export MODEL_NAME="Dolphin"
-
-python api_server.py \
-    --hf_model_dir tmp/hf_models/${MODEL_NAME} \
-    --visual_engine_dir tmp/trt_engines/${MODEL_NAME}/vision_encoder \
-    --llm_engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 \
-    --max_batch_size 16
-
-# 2. Predict
-# predict elements reading order
-python deployment/tensorrt_llm/api_client.py --image_path ./demo/page_imgs/page_1.jpeg --prompt "Parse the reading order of this document."
-
-# recognize text/latex
-python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/block_formula.jpeg --prompt "Read text in the image."
-python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/para_1.jpg --prompt "Read text in the image."
-
-# recognize table
-python deployment/tensorrt_llm/api_client.py --image_path ./demo/element_imgs/table_1.jpeg --prompt "Parse the table in the image."
-```

deployment/tensorrt_llm/api_client.py

@@ -1,100 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Example Python client for `vllm.entrypoints.api_server`
-Start the demo server:
-    python -m vllm.entrypoints.api_server --model <model_name>
-
-NOTE: The API server is used only for demonstration and simple performance
-benchmarks. It is not intended for production use.
-For production use, we recommend `vllm serve` and the OpenAI client API.
-"""
-
-import argparse
-import base64
-import json
-from argparse import Namespace
-from collections.abc import Iterable
-
-import requests
-
-
-def clear_line(n: int = 1) -> None:
-    LINE_UP = "\033[1A"
-    LINE_CLEAR = "\x1b[2K"
-    for _ in range(n):
-        print(LINE_UP, end=LINE_CLEAR, flush=True)
-
-
-def encode_image_base64(image_path: str) -> str:
-    """Encode local image to base64 format."""
-
-    with open(image_path, "rb") as f:
-        image_data = f.read()
-        result = base64.b64encode(image_data).decode("utf-8")
-
-    return result
-
-
-def post_http_request(
-        prompt: str, image_path: str, api_url: str, stream: bool = False
-) -> requests.Response:
-    headers = {"User-Agent": "Test Client"}
-    pload = {
-        "prompt": prompt,
-        "image_base64": encode_image_base64(image_path),
-    }
-    response = requests.post(api_url, headers=headers, json=pload, stream=stream)
-    return response
-
-
-def get_streaming_response(response: requests.Response) -> Iterable[list[str]]:
-    for chunk in response.iter_lines(
-            chunk_size=8192, decode_unicode=False, delimiter=b"\n"
-    ):
-        if chunk:
-            data = json.loads(chunk.decode("utf-8"))
-            output = data["text"]
-            yield output
-
-
-def get_response(response: requests.Response) -> list[str]:
-    data = json.loads(response.content)
-    output = data["text"]
-    return output
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", type=str, default="localhost")
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--prompt", type=str, default="Parse the reading order of this document.")
-    parser.add_argument("--image_path", type=str, default="./demo/page_imgs/page_1.jpeg")
-    parser.add_argument("--stream", action="store_true")
-    return parser.parse_args()
-
-
-def main(args: Namespace):
-    prompt = args.prompt
-    image_path = args.image_path
-    api_url = f"http://{args.host}:{args.port}/generate"
-    stream = args.stream
-
-    print(f"Prompt: {prompt!r}\n", flush=True)
-    response = post_http_request(prompt, image_path, api_url, stream)
-
-    if stream:
-        num_printed_lines = 0
-        for h in get_streaming_response(response):
-            clear_line(num_printed_lines)
-            num_printed_lines = 0
-            for i, line in enumerate(h):
-                num_printed_lines += 1
-                print(f"Response {i}: {line!r}", flush=True)
-    else:
-        output = get_response(response)
-        print(f"Response: {output!r}", flush=True)
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

deployment/tensorrt_llm/api_server.py

@@ -1,112 +0,0 @@
-# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.1/examples/apps/fastapi_server.py
-
-#!/usr/bin/env python
-import asyncio
-import base64
-import io
-import logging
-import signal
-from http import HTTPStatus
-from PIL import Image
-from typing import Optional
-
-import click
-import uvicorn
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse, Response
-
-from tensorrt_llm.executor import CppExecutorError, RequestError
-from dolphin_runner import DolphinRunner, InferenceConfig
-
-TIMEOUT_KEEP_ALIVE = 5  # seconds.
-
-
-async def decode_image(image_base64: str) -> Image.Image:
-    image_data = base64.b64decode(image_base64)
-    image = Image.open(io.BytesIO(image_data))
-    return image
-
-
-class LlmServer:
-    def __init__(self, runner: DolphinRunner):
-        self.runner = runner
-        self.app = FastAPI()
-        self.register_routes()
-
-    def register_routes(self):
-        self.app.add_api_route("/health", self.health, methods=["GET"])
-        self.app.add_api_route("/generate", self.generate, methods=["POST"])
-
-    async def health(self) -> Response:
-        return Response(status_code=200)
-
-    async def generate(self, request: Request) -> Response:
-        """ Generate completion for the request.
-
-        The request should be a JSON object with the following fields:
-        - prompt: the prompt to use for the generation.
-        - image_base64: the image to use for the generation.
-        """
-        request_dict = await request.json()
-
-        prompt = request_dict.pop("prompt", "")
-        logging.info(f"request prompt: {prompt}")
-        image_base64 = request_dict.pop("image_base64", "")
-        image = await decode_image(image_base64)
-
-        try:
-            output_texts = self.runner.run([prompt], [image], 4024)
-            output_texts = [texts[0] for texts in output_texts]
-            return JSONResponse({"text": output_texts[0]})
-        except RequestError as e:
-            return JSONResponse(content=str(e),
-                                status_code=HTTPStatus.BAD_REQUEST)
-        except CppExecutorError:
-            # If internal executor error is raised, shutdown the server
-            signal.raise_signal(signal.SIGINT)
-
-    async def __call__(self, host, port):
-        config = uvicorn.Config(self.app,
-                                host=host,
-                                port=port,
-                                log_level="info",
-                                timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
-        await uvicorn.Server(config).serve()
-
-
-@click.command()
-@click.option("--hf_model_dir", type=str, required=True)
-@click.option("--visual_engine_dir", type=str, required=True)
-@click.option("--llm_engine_dir", type=str, required=True)
-@click.option("--max_batch_size", type=int, default=16)
-@click.option("--max_new_tokens", type=int, default=4024)
-@click.option("--host", type=str, default=None)
-@click.option("--port", type=int, default=8000)
-def entrypoint(hf_model_dir: str,
-               visual_engine_dir: str,
-               llm_engine_dir: str,
-               max_batch_size: int,
-               max_new_tokens: int,
-               host: Optional[str] = None,
-               port: int = 8000):
-    host = host or "0.0.0.0"
-    port = port or 8000
-    logging.info(f"Starting server at {host}:{port}")
-
-    config = InferenceConfig(
-        max_new_tokens=max_new_tokens,
-        batch_size=max_batch_size,
-        log_level="info",
-        hf_model_dir=hf_model_dir,
-        visual_engine_dir=visual_engine_dir,
-        llm_engine_dir=llm_engine_dir,
-    )
-
-    dolphin_runner = DolphinRunner(config)
-    server = LlmServer(runner=dolphin_runner)
-
-    asyncio.run(server(host, port))
-
-
-if __name__ == "__main__":
-    entrypoint()

deployment/tensorrt_llm/convert/build_visual_engine.py

@@ -1,14 +0,0 @@
-# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.2/examples/multimodal/build_visual_engine.py
-
-import argparse
-
-from tensorrt_llm.tools.multimodal_builder import (VisionEngineBuilder,
-                                                   add_multimodal_arguments)
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser = add_multimodal_arguments(parser)
-    args = parser.parse_args()
-
-    builder = VisionEngineBuilder(args)
-    builder.build()

deployment/tensorrt_llm/convert/convert_checkpoint.py

@@ -1,1528 +0,0 @@
-# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.1/examples/enc_dec/convert_checkpoint.py
-
-import argparse
-import configparser
-import copy
-import json
-import logging
-import os
-import types
-from ast import literal_eval
-from datetime import datetime
-from pathlib import Path
-
-import safetensors
-from helper import convert_weight_to_dtype, fuse_qkv_one_layer, reshape, split
-from transformers import (AutoModelForSeq2SeqLM, Blip2ForConditionalGeneration,
-                          MBartForConditionalGeneration,
-                          Pix2StructForConditionalGeneration,
-                          T5ForConditionalGeneration, VisionEncoderDecoderModel)
-
-from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType,
-                                     MLPType)
-from tensorrt_llm.models import PretrainedConfig
-
-dir_path = os.path.dirname(os.path.realpath(__file__))
-LOGGER = logging.getLogger(__name__)
-
-layernorm_type_map = {i.name: i.value for i in LayerNormType}
-layernorm_position_map = {i.name: i.value for i in LayerNormPositionType}
-mlp_type_map = {i.name: i.value for i in MLPType}
-
-
-def copy_args_to_component_config(component_config, args):
-    for arg in vars(args):
-        setattr(component_config, arg, getattr(args, arg))
-    return component_config
-
-
-def parse_t5_config(args, hf_model):
-    config = configparser.ConfigParser()
-
-    config["encoder"] = {}
-    for key, val in hf_model.encoder.config.to_dict().items():
-        config["encoder"][key] = f"{val}"
-
-    # manually set q_scaling to offset attention scaling's effect.
-    # TODO: modify kernels to control whether to disable attention scaling
-    def get_offset_q_scaling(config):
-        scaling = 1 / config.head_size**.5
-        return scaling
-
-    config["decoder"] = {}
-    for key, val in hf_model.decoder.config.to_dict().items():
-        config["decoder"][key] = f"{val}"
-
-    config["structure"] = dict()
-    config["structure"]["t5_with_bias"] = "false"
-    config["structure"]["use_gated_activation"] = str(
-        hf_model.encoder.config.is_gated_act)
-    config["structure"]["position_embedding_type"] = "relative"
-    config["structure"]["model_type"] = args.model_type
-
-    def parse_t5_config_by_component(config, component, args):
-        component_config = types.SimpleNamespace()
-        component_config = copy_args_to_component_config(component_config, args)
-        component_config.n_head = config.getint(component, 'num_heads')
-        component_config.head_size = config.getint(component, 'd_kv')
-        component_config.hidden_size = config.getint(component, 'd_model')
-        component_config.ffn_hidden_size = config.getint(component, 'd_ff')
-        component_config.vocab_size = config.getint(component, 'vocab_size')
-        component_config.n_positions = config.getint(component,
-                                                     'n_positions',
-                                                     fallback=512)
-        component_config.has_position_embedding = config.getboolean(
-            component, 'has_position_embedding',
-            fallback=False)  # TODO: hardcoded here
-
-        component_config.has_token_type_embedding = config.getboolean(
-            component, 'has_token_type_embedding', fallback=False)
-        component_config.has_embedding_layernorm = config.getboolean(
-            component, 'has_embedding_layernorm', fallback=False)
-        component_config.has_embedding_scale = config.getboolean(
-            component, 'has_embedding_scale', fallback=False)
-        component_config.q_scaling = get_offset_q_scaling(component_config)
-        component_config.has_attention_qkvo_bias = config.getboolean(
-            component, 'has_attention_qkvo_bias',
-            fallback=False)  # TODO: hardcoded here
-        component_config.has_mlp_bias = config.getboolean(component,
-                                                          'has_mlp_bias',
-                                                          fallback=False)
-        component_config.has_model_final_layernorm = config.getboolean(
-            component, 'has_model_final_layernorm', fallback=True)
-        component_config.layernorm_eps = config.getfloat(
-            component, 'layer_norm_epsilon')
-        component_config.layernorm_position = layernorm_position_map[config.get(
-            component, 'layernorm_position',
-            fallback='pre_layernorm')]  # TODO: hardcoded here
-        component_config.layernorm_type = layernorm_type_map[config.get(
-            component, 'layernorm_type', fallback='RmsNorm')]
-        component_config.hidden_act = config.get(component, 'dense_act_fn')
-        component_config.gated_act = config.getboolean(component,
-                                                       'is_gated_act')
-        component_config.mlp_type = mlp_type_map['GatedMLP' if component_config.
-                                                 gated_act else 'MLP']
-        component_config.num_buckets = config.getint(
-            component, 'relative_attention_num_buckets')
-        component_config.max_distance = config.getint(
-            component, 'relative_attention_max_distance')
-        component_config.position_embedding_type = config.get(
-            'structure', 'position_embedding_type')
-        component_config.logits_dtype = config.get(component,
-                                                   'logits_dtype',
-                                                   fallback='float32')
-
-        if component == 'encoder':
-            component_config.n_layer = config.getint(component, 'num_layers')
-
-            component_config.relative_attention = config.get(
-                'structure', 'position_embedding_type') == 'relative'
-
-        elif component == 'decoder':
-            component_config.n_layer = config.getint(component,
-                                                     'num_decoder_layers')
-            component_config.has_lm_head_bias = config.getboolean(
-                component,  # TODO: T5 with bias
-                'has_lm_head_bias',
-                fallback=False)
-            component_config.relative_attention = config.getboolean(
-                component, 'relative_attention', fallback=True)
-            component_config.rescale_before_lm_head = config.getboolean(
-                component, 'tie_word_embeddings'
-            )  # default is True (for T5), but False for Flan-T5
-            component_config.encoder_hidden_size = config.getint(
-                'encoder', 'd_model')
-            component_config.encoder_num_heads = config.getint(
-                'encoder', 'num_heads')
-            component_config.encoder_head_size = config.getint(
-                'encoder', 'd_kv')
-            component_config.decoder_start_token_id = config.getint(
-                'decoder', 'decoder_start_token_id')
-            component_config.eos_token_id = config.getint(
-                'decoder', 'eos_token_id')
-            bos_token_id = config.get('decoder', 'bos_token_id')
-            # T5 does not have bos_token_id
-            component_config.bos_token_id = int(
-                bos_token_id) if bos_token_id != "None" else None
-            component_config.pad_token_id = config.getint(
-                'decoder', 'pad_token_id')
-
-        else:
-            assert False, 'Unsupported component!'
-
-        return component_config
-
-    encoder_config = parse_t5_config_by_component(config, "encoder", args)
-    decoder_config = parse_t5_config_by_component(config, "decoder", args)
-
-    return encoder_config, decoder_config
-
-
-def convert_t5_weights_to_tllm_safetensors(config, component, params):
-    weights = {}
-
-    mapping = config.mapping
-
-    convert_weight_to_dtype(params, config.dtype)
-    hidden_size = config.hidden_size
-    ffn_hidden_size = config.intermediate_size
-    num_layers = config.num_hidden_layers
-    n_head = config.num_attention_heads
-    head_size = config.head_size
-    attention_hidden_size = n_head * head_size  # head size * num_heads not necessarily equals hidden_dim, such as Flan-T5
-
-    hf_param_prefix = f'{component}'
-    trtllm_layer_name = f'{component}_layers'
-    trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention'
-    trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm'
-    hf_component_idx = 1 if component == 'encoder' else 2
-
-    def get_attn_module_name(component, block, layer, attn_type):
-        return f'{component}.block.{int(block)}.layer.{int(layer)}.{attn_type}'
-
-    weights['embedding.vocab_embedding.weight'] = reshape(
-        params['shared.weight'].clone(), None)
-
-    layers_range = mapping.pp_layers(num_layers)
-    for layer_idx in layers_range:
-        local_layer_idx = layer_idx - layers_range[0]
-        trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}'
-        hf_layer_name_prefix = f'{hf_param_prefix}.block.{layer_idx}'
-
-        hidden_layer_name_split = {
-            f'{hf_layer_name_prefix}.layer.0.SelfAttention.o.weight': {
-                "name":
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.dense.weight',
-                "shape":
-                (hidden_size, attention_hidden_size // mapping.tp_size),
-                "split_dim": -1
-            },
-            f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wo.weight':
-            {
-                "name": f'{trtllm_layer_name_prefix}.mlp.proj.weight',
-                "shape": (hidden_size, ffn_hidden_size // mapping.tp_size),
-                "split_dim": -1
-            },
-            f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi.weight':
-            {
-                "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight',
-                "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-                "split_dim": 0
-            },
-            f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi_0.weight':
-            {
-                "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight',
-                "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-                "split_dim": 0
-            },
-        }
-
-        hidden_layer_name_no_split = {
-            f'{hf_layer_name_prefix}.layer.0.layer_norm.weight': {
-                "name":
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layernorm_name}.weight',
-                "shape": None
-            },
-            f'{hf_layer_name_prefix}.layer.{hf_component_idx}.layer_norm.weight':
-            {
-                "name": f'{trtllm_layer_name_prefix}.mlp_layernorm.weight',
-                "shape": None
-            },
-        }
-
-        if config.gated_act:
-            hidden_layer_name_split.update({
-                f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi2.weight':
-                {
-                    "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight',
-                    "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-                    "split_dim": 0
-                },
-                f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi_1.weight':
-                {
-                    "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight',
-                    "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-                    "split_dim": 0
-                },
-            })
-
-        if component == 'decoder':
-            hidden_layer_name_split.update({
-                f'{hf_layer_name_prefix}.layer.1.EncDecAttention.o.weight': {
-                    "name":
-                    f'{trtllm_layer_name_prefix}.cross_attention.dense.weight',
-                    "shape":
-                    (hidden_size, attention_hidden_size // mapping.tp_size),
-                    "split_dim": -1
-                },
-            })
-            hidden_layer_name_no_split.update({
-                f'{hf_layer_name_prefix}.layer.1.layer_norm.weight': {
-                    "name":
-                    f'{trtllm_layer_name_prefix}.cross_attention_layernorm.weight',
-                    "shape": None
-                },
-            })
-            self_attn_module_name = get_attn_module_name(
-                component, layer_idx, "1", 'EncDecAttention')
-            weights.update(
-                fuse_qkv_one_layer(
-                    params, self_attn_module_name,
-                    f'{trtllm_layer_name_prefix}.cross_attention',
-                    mapping.tp_size, mapping.tp_rank, config.model_type,
-                    (attention_hidden_size * 3 // mapping.tp_size, hidden_size),
-                    None))
-
-        self_attn_module_name = get_attn_module_name(component, layer_idx, "0",
-                                                     'SelfAttention')
-        weights.update(
-            fuse_qkv_one_layer(
-                params, self_attn_module_name,
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}',
-                mapping.tp_size, mapping.tp_rank, config.model_type,
-                (attention_hidden_size * 3 // mapping.tp_size, hidden_size),
-                None))
-
-        weights[
-            f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.rel_attn_table'] = reshape(
-                split(
-                    params[
-                        f'{component}.block.0.layer.0.SelfAttention.relative_attention_bias.weight']
-                    .T, mapping.tp_size, mapping.tp_rank, 0),
-                (n_head // mapping.tp_size, config.num_buckets))
-
-        for hf_weight_name, weight_info in hidden_layer_name_split.items():
-            if hf_weight_name in params.keys():
-                weights[weight_info["name"]] = reshape(
-                    split(params[hf_weight_name],
-                          mapping.tp_size,
-                          mapping.tp_rank,
-                          dim=weight_info["split_dim"]), weight_info["shape"])
-        for hf_weight_name, weight_info in hidden_layer_name_no_split.items():
-            if hf_weight_name in params.keys():
-                weights[weight_info["name"]] = reshape(
-                    params[hf_weight_name].clone(), shape=weight_info["shape"])
-
-    weights['final_layernorm.weight'] = reshape(
-        params[f'{component}.final_layer_norm.weight'].clone(), None)
-
-    if component == 'decoder':
-        weights['lm_head.weight'] = reshape(
-            split(params['lm_head.weight'],
-                  mapping.tp_size,
-                  mapping.tp_rank,
-                  dim=0), (config.vocab_size // mapping.tp_size, hidden_size))
-        if not config.use_implicit_relative_attention:
-            weights['rel_attn_table'] = reshape(
-                split(
-                    params[
-                        f'{component}.block.0.layer.0.SelfAttention.relative_attention_bias.weight']
-                    .T, mapping.tp_size, mapping.tp_rank, 0),
-                (n_head // mapping.tp_size, config.num_buckets))
-
-    return weights
-
-
-convert_blip2_weights_to_tllm_safetensors = convert_t5_weights_to_tllm_safetensors  # func alias
-
-
-def parse_nmt_config(args, model):
-    config = configparser.ConfigParser()
-    fairseq_config = vars(model.cfg.model)  # Namespace --> dict
-
-    config['encoder'] = dict()
-    for key, val in fairseq_config.items():
-        config["encoder"][key] = f"{val}"
-    config["encoder"]["q_scaling"] = '1'
-    # NMT has final layernorm for pre-norm model architecture.
-    config['encoder']['has_model_final_layernorm'] = config['encoder'][
-        'encoder_normalize_before']
-    config['encoder']['vocab_size'] = str(len(model.src_dict))  # fairseq naming
-
-    config['decoder'] = dict()
-    for key, val in fairseq_config.items():
-        config["decoder"][key] = f"{val}"
-    config["decoder"]["q_scaling"] = '1'
-    config["decoder"]["rescale_before_lm_head"] = 'false'
-    config['decoder']['has_model_final_layernorm'] = str(
-        config['decoder'].getboolean('decoder_normalize_before', False)
-        and not config['decoder'].getboolean('no_decoder_final_norm', False))
-    config['decoder']['vocab_size'] = str(len(model.tgt_dict))  # fairseq naming
-
-    config["structure"] = dict()
-    config["structure"]["t5_with_bias"] = "true"
-    config["structure"]["use_gated_activation"] = "false"
-    config["structure"][
-        "position_embedding_type"] = "learned_absolute"  # "sinusoid"
-    config["structure"]["model_type"] = args.model_type
-
-    def parse_nmt_config_by_component(config, component, args):
-        assert component in ('encoder', 'decoder'), 'Unsupported component!'
-        component_config = types.SimpleNamespace()
-        component_config = copy_args_to_component_config(component_config, args)
-        component_config.n_layer = config.getint(component,
-                                                 f'{component}_layers')
-        component_config.n_head = config.getint(component,
-                                                f'{component}_attention_heads')
-        component_config.hidden_size = config.getint(
-            component, f'{component}_embed_dim')  # fairseq naming
-        component_config.head_size = config.getint(
-            component,
-            'd_kv',
-            fallback=component_config.hidden_size // component_config.n_head)
-        component_config.ffn_hidden_size = config.getint(
-            component, f'{component}_ffn_embed_dim')  # fairseq naming
-        component_config.vocab_size = config.getint(component, 'vocab_size')
-        component_config.n_positions = config.getint(
-            component, 'max_source_positions')  # fairseq naming
-        component_config.has_position_embedding = not config.getboolean(
-            component, 'no_token_positional_embeddings',
-            fallback=False)  # fairseq naming
-        component_config.has_token_type_embedding = config.getboolean(
-            component, 'has_token_type_embedding', fallback=False)
-        component_config.has_embedding_layernorm = config.getboolean(
-            component, 'layernorm_embedding', fallback=True)  # fairseq naming
-        component_config.has_embedding_scale = not config.getboolean(
-            component, 'no_scale_embedding')  # fairseq naming
-        component_config.q_scaling = config.getfloat(component,
-                                                     'q_scaling',
-                                                     fallback=1.0)
-        component_config.has_attention_qkvo_bias = config.getboolean(
-            'structure', 't5_with_bias', fallback=True)
-        component_config.has_mlp_bias = config.getboolean('structure',
-                                                          't5_with_bias',
-                                                          fallback=True)
-        component_config.has_model_final_layernorm = config.getboolean(
-            component, 'has_model_final_layernorm')
-        component_config.layernorm_eps = config.getfloat(
-            component, 'layer_norm_epsilon', fallback=1e-5)  # fairseq naming
-
-        normalize_before = config.getboolean(
-            component, f'{component}_normalize_before')  # fairseq naming
-        component_config.layernorm_position = layernorm_position_map[
-            'pre_layernorm' if normalize_before else 'post_layernorm']
-
-        component_config.layernorm_type = layernorm_type_map[config.get(
-            component, 'layernorm_type', fallback='LayerNorm')]
-        component_config.hidden_act = config.get(
-            component, 'activation_fn')  # fairseq naming
-        component_config.gated_act = config.getboolean(component,
-                                                       'is_gated_act',
-                                                       fallback=False)
-        component_config.mlp_type = mlp_type_map['GatedMLP' if component_config.
-                                                 gated_act else 'MLP']
-        component_config.relative_attention = config.get(
-            'structure', 'position_embedding_type') == 'relative'
-
-        component_config.num_buckets = config.getint(
-            component, 'relative_attention_num_buckets', fallback=0)
-        component_config.max_distance = config.getint(
-            component, 'relative_attention_max_distance', fallback=0)
-        component_config.position_embedding_type = config.get(
-            'structure', 'position_embedding_type')
-        component_config.logits_dtype = config.get(component,
-                                                   'logits_dtype',
-                                                   fallback='float32')
-        if component == 'decoder':
-            component_config.rescale_before_lm_head = config.getboolean(
-                component, 'rescale_before_lm_head')
-
-            component_config.encoder_hidden_size = config.getint(
-                'encoder', 'encoder_embed_dim')  # fairseq naming
-            component_config.encoder_num_heads = config.getint(
-                'encoder', 'encoder_attention_heads')
-            component_config.encoder_head_size = config.getint(
-                'encoder',
-                'd_kv',
-                fallback=component_config.encoder_hidden_size //
-                component_config.encoder_num_heads)
-            component_config.decoder_start_token_id = None
-            component_config.eos_token_id = None
-            component_config.bos_token_id = None
-            component_config.pad_token_id = None
-
-        return component_config
-
-    encoder_config = parse_nmt_config_by_component(config, "encoder", args)
-    decoder_config = parse_nmt_config_by_component(config, "decoder", args)
-
-    return encoder_config, decoder_config
-
-
-def convert_nmt_weights_to_tllm_safetensors(config, component, params,
-                                            sin_pos_embedding):
-    weights = {}
-
-    mapping = config.mapping
-
-    hidden_size = config.hidden_size
-
-    convert_weight_to_dtype(params, config.dtype)
-    ffn_hidden_size = config.intermediate_size
-    vocab_size = config.vocab_size
-
-    hf_param_prefix = f'models.0.{component}'
-    trtllm_layer_name = f'{component}_layers'
-    trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention'
-    trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm'
-
-    hidden_layer_name_split = {
-        'self_attn.out_proj.weight': {
-            "name": f'{trtllm_attn_layer_name}.dense.weight',
-            "shape": (hidden_size, hidden_size // mapping.tp_size),
-            "split_dim": -1
-        },
-        'fc1.weight': {
-            "name": 'mlp.fc.weight',
-            "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-            "split_dim": 0
-        },
-        'fc1.bias': {
-            "name": 'mlp.fc.bias',
-            "shape": (ffn_hidden_size // mapping.tp_size),
-            "split_dim": 0
-        },
-        'fc2.weight': {
-            "name": 'mlp.proj.weight',
-            "shape": (hidden_size, ffn_hidden_size // mapping.tp_size),
-            "split_dim": -1
-        },
-    }
-
-    hidden_layer_name_no_split = {
-        'self_attn.out_proj.bias': {
-            "name": f'{trtllm_attn_layer_name}.dense.bias',
-            "shape": (hidden_size)
-        },
-        'self_attn_layer_norm.weight': {
-            "name": f'{trtllm_attn_layernorm_name}.weight',
-            "shape": None
-        },
-        'self_attn_layer_norm.bias': {
-            "name": f'{trtllm_attn_layernorm_name}.bias',
-            "shape": None
-        },
-        'fc2.bias': {
-            "name": 'mlp.proj.bias',
-            "shape": (hidden_size)
-        },
-        'final_layer_norm.weight': {
-            "name": 'mlp_layernorm.weight',
-            "shape": None
-        },
-        'final_layer_norm.bias': {
-            "name": 'mlp_layernorm.bias',
-            "shape": None
-        },
-    }
-
-    if component == "decoder":
-        hidden_layer_name_split.update({
-            'encoder_attn.out_proj.weight': {
-                "name": 'cross_attention.dense.weight',
-                "shape": (hidden_size, hidden_size // mapping.tp_size),
-                "split_dim": -1
-            },
-        })
-        hidden_layer_name_no_split.update({
-            'encoder_attn.out_proj.bias': {
-                "name": 'cross_attention.dense.bias',
-                "shape": (hidden_size)
-            },
-            'encoder_attn_layer_norm.weight': {
-                "name": 'cross_attention_layernorm.weight',
-                "shape": None,
-            },
-            'encoder_attn_layer_norm.bias': {
-                "name": 'cross_attention_layernorm.bias',
-                "shape": None
-            },
-        })
-
-    def get_attn_module_name(component, layer, attn_type):
-        return f'models.0.{component}.layers.{int(layer)}.{attn_type}'
-
-    weights["embedding.vocab_embedding.weight"] = reshape(
-        params[f'{hf_param_prefix}.embed_tokens.weight'].clone(),
-        (vocab_size, -1))
-    weights["embedding.position_embedding.weight"] = reshape(
-        sin_pos_embedding, (config.max_position_embeddings, hidden_size))
-
-    num_layers = config.num_hidden_layers
-
-    layers_range = mapping.pp_layers(num_layers)
-    for layer_idx in layers_range:
-        local_layer_idx = layer_idx - layers_range[0]
-        hf_layer_name_prefix = f'{hf_param_prefix}.layers.{layer_idx}'
-        trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}'
-
-        for hf_weight_name, weight_info in hidden_layer_name_split.items():
-            weights[
-                f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape(
-                    split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'],
-                          mapping.tp_size,
-                          mapping.tp_rank,
-                          dim=weight_info["split_dim"]), weight_info["shape"])
-
-        for hf_weight_name, weight_info in hidden_layer_name_no_split.items():
-            trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}'
-            hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}'
-            weights[trtllm_layer_fullname] = reshape(
-                params[hf_layer_fullname].clone(), shape=weight_info["shape"])
-
-        self_attn_module_name = get_attn_module_name(component, layer_idx,
-                                                     'self_attn')
-        weights.update(
-            fuse_qkv_one_layer(
-                params, self_attn_module_name,
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}',
-                mapping.tp_size, mapping.tp_rank, config.model_type,
-                (hidden_size * 3 // mapping.tp_size, hidden_size),
-                (hidden_size * 3 // mapping.tp_size)))
-        if component == 'decoder':
-            cross_attn_module_name = get_attn_module_name(
-                component, layer_idx, 'encoder_attn')
-            weights.update(
-                fuse_qkv_one_layer(
-                    params, cross_attn_module_name,
-                    f'{trtllm_layer_name_prefix}.cross_attention',
-                    mapping.tp_size, mapping.tp_rank, config.model_type,
-                    (hidden_size * 3 // mapping.tp_size, hidden_size),
-                    (hidden_size * 3 // mapping.tp_size)))
-
-    if component == 'decoder':
-        weights['lm_head.weight'] = reshape(
-            split(params[f'{hf_param_prefix}.output_projection.weight'],
-                  mapping.tp_size,
-                  mapping.tp_rank,
-                  dim=0), (config.vocab_size // mapping.tp_size, hidden_size))
-
-    if config.has_model_final_layernorm:
-        weights['final_layernorm.weight'] = params[
-            f'{hf_param_prefix}.layer_norm.weight'].clone()
-        weights['final_layernorm.bias'] = params[
-            f'{hf_param_prefix}.layer_norm.bias'].clone()
-
-    return weights
-
-
-def parse_bart_config(args, hf_model):
-
-    config = configparser.ConfigParser()
-
-    config['decoder'] = dict()
-    for key, val in hf_model.model.decoder.config.to_dict().items():
-        config["decoder"][key] = f"{val}"
-    config["decoder"]["q_scaling"] = '1'
-    config["decoder"]["rescale_before_lm_head"] = str(False)
-    config['decoder']['has_model_final_layernorm'] = str(
-        args.nougat or isinstance(hf_model, MBartForConditionalGeneration))
-
-    if args.nougat:
-        # These flags are true for mbart decoders, but missing in HF config
-        config['decoder']['normalize_before'] = str(True)
-        config['decoder']['normalize_embeddings'] = str(True)
-
-        config['encoder'] = dict()
-        # Init few encoder configs, needed by build, from decoder config
-        encoder_config_keys = [
-            "encoder_ffn_dim", "encoder_layers", "encoder_attention_heads",
-            "encoder_layerdrop", "d_model"
-        ]
-        for key in encoder_config_keys:
-            config['encoder'][key] = config['decoder'][key]
-    else:
-        config['encoder'] = dict()
-        for key, val in hf_model.model.encoder.config.to_dict().items():
-            config["encoder"][key] = f"{val}"
-        config["encoder"]["q_scaling"] = '1'
-
-        # mBART has final layernorm, BART does not
-        config['encoder']['has_model_final_layernorm'] = str(
-            isinstance(hf_model, MBartForConditionalGeneration))
-
-    config["structure"] = dict()
-    config["structure"]["t5_with_bias"] = "true"
-    config["structure"]["use_gated_activation"] = "false"
-    config["structure"]["position_embedding_type"] = "learned_absolute"
-    config["structure"]["model_type"] = args.model_type
-
-    def parse_bart_config_by_component(config, component, args):
-        assert component in ('encoder', 'decoder'), 'Unsupported component!'
-        component_config = types.SimpleNamespace()
-        component_config = copy_args_to_component_config(component_config, args)
-        component_config.n_layer = config.getint(component,
-                                                 f'{component}_layers')
-        component_config.n_head = config.getint(component,
-                                                f'{component}_attention_heads')
-        component_config.hidden_size = config.getint(component, 'd_model')
-        component_config.head_size = config.getint(
-            component,
-            'd_kv',
-            fallback=component_config.hidden_size // component_config.n_head)
-        component_config.ffn_hidden_size = config.getint(
-            component, f'{component}_ffn_dim')
-        component_config.vocab_size = config.getint(component, 'vocab_size')
-        component_config.n_positions = config.getint(component,
-                                                     'max_position_embeddings')
-        component_config.has_position_embedding = config.getboolean(
-            component, 'has_position_embedding',
-            fallback=True)  # TODO: hardcoded here
-        component_config.has_token_type_embedding = config.getboolean(
-            component, 'has_token_type_embedding', fallback=False)
-        component_config.has_embedding_layernorm = config.getboolean(
-            component, 'has_embedding_layernorm', fallback=True)
-        component_config.has_embedding_scale = config.getboolean(
-            component, 'scale_embedding')
-        component_config.q_scaling = config.getfloat(component,
-                                                     'q_scaling',
-                                                     fallback=1.0)
-        component_config.has_attention_qkvo_bias = config.getboolean(
-            'structure', 't5_with_bias', fallback=True)
-        component_config.has_mlp_bias = config.getboolean('structure',
-                                                          't5_with_bias',
-                                                          fallback=True)
-        component_config.has_model_final_layernorm = config.getboolean(
-            component, 'has_model_final_layernorm')
-        component_config.layernorm_eps = config.getfloat(component,
-                                                         'layer_norm_epsilon',
-                                                         fallback=False)
-
-        normalize_before = config.getboolean(component, 'normalize_before')
-        component_config.layernorm_position = layernorm_position_map[
-            'pre_layernorm' if normalize_before else 'post_layernorm']
-
-        component_config.layernorm_type = layernorm_type_map[config.get(
-            component, 'layernorm_type', fallback='LayerNorm')]
-        component_config.hidden_act = config.get(component,
-                                                 'activation_function')
-        component_config.gated_act = config.getboolean(component,
-                                                       'is_gated_act',
-                                                       fallback=False)
-        component_config.mlp_type = mlp_type_map['GatedMLP' if component_config.
-                                                 gated_act else 'MLP']
-        component_config.relative_attention = config.get(
-            'structure', 'position_embedding_type') == 'relative'
-
-        component_config.num_buckets = config.getint(
-            component, 'relative_attention_num_buckets', fallback=0)
-        component_config.max_distance = config.getint(
-            component, 'relative_attention_max_distance', fallback=0)
-        component_config.max_lora_rank = config.getint(component,
-                                                       'max_lora_rank',
-                                                       fallback=0)
-        component_config.lora_target_modules = literal_eval(
-            config.get(component, 'lora_target_modules', fallback="[]"))
-        component_config.hf_modules_to_trtllm_modules = literal_eval(
-            config.get(component, 'hf_modules_to_trtllm_modules',
-                       fallback="{}"))
-        component_config.trtllm_modules_to_hf_modules = literal_eval(
-            config.get(component, 'trtllm_modules_to_hf_modules',
-                       fallback="{}"))
-        component_config.logits_dtype = config.get(component,
-                                                   'logits_dtype',
-                                                   fallback='float32')
-        component_config.position_embedding_type = config.get(
-            'structure', 'position_embedding_type')
-
-        if component == 'decoder':
-            component_config.rescale_before_lm_head = config.getboolean(
-                component, 'rescale_before_lm_head')
-
-            component_config.encoder_hidden_size = config.getint(
-                'encoder', 'd_model')
-            component_config.encoder_num_heads = config.getint(
-                'encoder', 'encoder_attention_heads')
-            component_config.encoder_head_size = config.getint(
-                'encoder',
-                'd_kv',
-                fallback=component_config.encoder_hidden_size //
-                component_config.encoder_num_heads)
-
-            # nougat has decoder_start_token_id = None, special handling
-            decoder_start_token_id = config.get('decoder',
-                                                'decoder_start_token_id')
-            component_config.decoder_start_token_id = int(
-                decoder_start_token_id
-            ) if decoder_start_token_id != "None" else None
-            component_config.eos_token_id = config.getint(
-                'decoder', 'eos_token_id')
-            component_config.bos_token_id = config.getint(
-                'decoder', 'bos_token_id')
-            component_config.pad_token_id = config.getint(
-                'decoder', 'pad_token_id')
-
-        return component_config
-
-    encoder_config = None
-    if not args.nougat:
-        encoder_config = parse_bart_config_by_component(config, "encoder", args)
-    decoder_config = parse_bart_config_by_component(config, "decoder", args)
-
-    return encoder_config, decoder_config
-
-
-def convert_bart_weights_to_tllm_safetensors(config, component, params):
-    weights = {}
-
-    mapping = config.mapping
-
-    hidden_size = config.hidden_size
-
-    convert_weight_to_dtype(params, config.dtype)
-    ffn_hidden_size = config.intermediate_size
-    vocab_size = config.vocab_size
-
-    hf_param_prefix = f'model.{component}'
-    trtllm_layer_name = f'{component}_layers'
-    trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention'
-    trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm'
-    embedding_layer_names = {
-        'embed_tokens.weight': {
-            "name": 'embedding.vocab_embedding.weight',
-            "shape": (vocab_size, -1)
-        },
-        'embed_positions.weight': {
-            "name": 'embedding.position_embedding.weight',
-            "shape": (config.max_position_embeddings, hidden_size)
-        },
-        'layernorm_embedding.weight': {
-            "name": 'embedding.embedding_layernorm.weight',
-            "shape": None
-        },
-        'layernorm_embedding.bias': {
-            "name": 'embedding.embedding_layernorm.bias',
-            "shape": None
-        },
-    }
-
-    hidden_layer_name_split = {
-        'self_attn.out_proj.weight': {
-            "name": f'{trtllm_attn_layer_name}.dense.weight',
-            "shape": (hidden_size, hidden_size // mapping.tp_size),
-            "split_dim": -1
-        },
-        'fc1.weight': {
-            "name": 'mlp.fc.weight',
-            "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-            "split_dim": 0
-        },
-        'fc1.bias': {
-            "name": 'mlp.fc.bias',
-            "shape": (ffn_hidden_size // mapping.tp_size),
-            "split_dim": 0
-        },
-        'fc2.weight': {
-            "name": 'mlp.proj.weight',
-            "shape": (hidden_size, ffn_hidden_size // mapping.tp_size),
-            "split_dim": -1
-        },
-    }
-
-    hidden_layer_name_no_split = {
-        'self_attn.out_proj.bias': {
-            "name": f'{trtllm_attn_layer_name}.dense.bias',
-            "shape": (hidden_size)
-        },
-        'self_attn_layer_norm.weight': {
-            "name": f'{trtllm_attn_layernorm_name}.weight',
-            "shape": None
-        },
-        'self_attn_layer_norm.bias': {
-            "name": f'{trtllm_attn_layernorm_name}.bias',
-            "shape": None
-        },
-        'fc2.bias': {
-            "name": 'mlp.proj.bias',
-            "shape": (hidden_size)
-        },
-        'final_layer_norm.weight': {
-            "name": 'mlp_layernorm.weight',
-            "shape": None
-        },
-        'final_layer_norm.bias': {
-            "name": 'mlp_layernorm.bias',
-            "shape": None
-        },
-    }
-
-    if config.model_type == 'mbart':
-        hidden_layer_name_split['layer_norm.weight'] = {
-            "name": 'final_layernorm.weight',
-            "shape": None,
-            "split_dim": 0
-        }
-        hidden_layer_name_no_split['layer_norm.bias'] = {
-            "name": 'final_layernorm.bias',
-            "shape": None,
-            "split_dim": 0
-        }
-
-    if component == "decoder":
-        hidden_layer_name_split.update({
-            'encoder_attn.out_proj.weight': {
-                "name": 'cross_attention.dense.weight',
-                "shape": (hidden_size, hidden_size // mapping.tp_size),
-                "split_dim": -1
-            }
-        })
-        hidden_layer_name_no_split.update({
-            'encoder_attn.out_proj.bias': {
-                "name": 'cross_attention.dense.bias',
-                "shape": (hidden_size)
-            },
-            'encoder_attn_layer_norm.weight': {
-                "name": 'cross_attention_layernorm.weight',
-                "shape": None
-            },
-            'encoder_attn_layer_norm.bias': {
-                "name": 'cross_attention_layernorm.bias',
-                "shape": None
-            },
-        })
-
-    def get_attn_module_name(component, layer, attn_type):
-        return f'model.{component}.layers.{int(layer)}.{attn_type}'
-
-    for hf_weight_name, weight_info in embedding_layer_names.items():
-        if 'position' in hf_weight_name:
-            weights[weight_info["name"]] = params[
-                f'{hf_param_prefix}.{hf_weight_name}'][2:].clone()
-        else:
-            weights[weight_info["name"]] = params[
-                f'{hf_param_prefix}.{hf_weight_name}'].clone()
-        weights[weight_info["name"]] = reshape(weights[weight_info["name"]],
-                                               weight_info["shape"])
-
-    num_layers = config.num_hidden_layers
-
-    layers_range = mapping.pp_layers(num_layers)
-    for layer_idx in layers_range:
-        local_layer_idx = layer_idx - layers_range[0]
-        hf_layer_name_prefix = f'{hf_param_prefix}.layers.{layer_idx}'
-        trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}'
-
-        for hf_weight_name, weight_info in hidden_layer_name_split.items():
-            weights[
-                f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape(
-                    split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'],
-                          mapping.tp_size,
-                          mapping.tp_rank,
-                          dim=weight_info["split_dim"]), weight_info["shape"])
-
-        for hf_weight_name, weight_info in hidden_layer_name_no_split.items():
-            trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}'
-            hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}'
-            weights[trtllm_layer_fullname] = reshape(
-                params[hf_layer_fullname].clone(), shape=weight_info["shape"])
-
-        self_attn_module_name = get_attn_module_name(component, layer_idx,
-                                                     'self_attn')
-        weights.update(
-            fuse_qkv_one_layer(
-                params, self_attn_module_name,
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}',
-                mapping.tp_size, mapping.tp_rank, config.model_type,
-                (hidden_size * 3 // mapping.tp_size, hidden_size),
-                (hidden_size * 3 // mapping.tp_size)))
-        if component == 'decoder':
-            cross_attn_module_name = get_attn_module_name(
-                component, layer_idx, 'encoder_attn')
-            weights.update(
-                fuse_qkv_one_layer(
-                    params, cross_attn_module_name,
-                    f'{trtllm_layer_name_prefix}.cross_attention',
-                    mapping.tp_size, mapping.tp_rank, config.model_type,
-                    (hidden_size * 3 // mapping.tp_size, hidden_size),
-                    (hidden_size * 3 // mapping.tp_size)))
-
-    if component == 'decoder':
-        weights['lm_head.weight'] = reshape(
-            split(params['lm_head.weight'],
-                  mapping.tp_size,
-                  mapping.tp_rank,
-                  dim=0), (config.vocab_size // mapping.tp_size, hidden_size))
-
-    if config.has_model_final_layernorm:
-        weights['final_layernorm.weight'] = params[
-            f'{hf_param_prefix}.layer_norm.weight'].clone()
-        weights['final_layernorm.bias'] = params[
-            f'{hf_param_prefix}.layer_norm.bias'].clone()
-
-    return weights
-
-
-def parse_pix2struct_config(args, hf_model):
-    # manually set q_scaling to offset attention scaling's effect.
-    # TODO: modify kernels to control whether to disable attention scaling
-    config = configparser.ConfigParser()
-
-    def get_offset_q_scaling(config) -> str:
-        d_model = config.hidden_size
-        num_heads = config.num_heads
-        head_size = d_model / num_heads
-        scaling = 1 / head_size**.5
-        return str(scaling)
-
-    config["decoder"] = {}
-    for key, val in hf_model.decoder.config.to_dict().items():
-        config["decoder"][key] = f"{val}"
-
-    config["decoder"]["q_scaling"] = get_offset_q_scaling(
-        hf_model.decoder.config)
-
-    config["structure"] = dict()
-    config["structure"]["pix2struct_with_bias"] = "false"
-    config["structure"]["use_gated_activation"] = "false"
-    config["structure"]["position_embedding_type"] = "relative"
-    config["structure"]["model_type"] = args.model_type
-
-    def parse_pix2struct_config_by_component(config, component, args):
-        if component == 'decoder':
-            args.n_layer = config.getint(component, 'num_layers')
-            args.n_head = config.getint(component, 'num_heads')
-            args.head_size = config.getint(component, 'd_kv')
-            args.hidden_size = config.getint(component, 'hidden_size')
-            args.ffn_hidden_size = config.getint(component, 'd_ff')
-            args.vocab_size = config.getint(component, 'vocab_size')
-            args.n_positions = config.getint(component,
-                                             'n_positions',
-                                             fallback=512)
-            args.has_position_embedding = config.getboolean(
-                component, 'has_position_embedding',
-                fallback=False)  # TODO: hardcoded here
-            args.has_token_type_embedding = config.getboolean(
-                component, 'has_token_type_embedding', fallback=False)
-            args.has_embedding_layernorm = config.getboolean(
-                component, 'has_embedding_layernorm', fallback=False)
-            args.has_embedding_scale = config.getboolean(component,
-                                                         'has_embedding_scale',
-                                                         fallback=False)
-            args.q_scaling = config.getfloat(component,
-                                             'q_scaling',
-                                             fallback=1.0)
-            args.has_attention_qkvo_bias = config.getboolean(
-                component, 'has_attention_qkvo_bias', fallback=False)
-            args.has_mlp_bias = config.getboolean(component,
-                                                  'has_mlp_bias',
-                                                  fallback=False)
-            args.has_model_final_layernorm = config.getboolean(
-                component, 'has_model_final_layernorm', fallback=True)
-            args.layernorm_eps = config.getfloat(component,
-                                                 'layer_norm_epsilon')
-            args.layernorm_position = layernorm_position_map[config.get(
-                component, 'layernorm_position',
-                fallback='pre_layernorm')]  # TODO: hardcoded here
-            args.layernorm_type = layernorm_type_map[config.get(
-                component, 'layernorm_type', fallback='RmsNorm')]
-            args.hidden_act = config.get(component, 'dense_act_fn')
-            args.gated_act = True
-            args.mlp_type = mlp_type_map['GatedMLP' if args.
-                                         gated_act else 'MLP']
-            args.has_lm_head_bias = config.getboolean(
-                component,  # TODO: T5 with bias
-                'has_lm_head_bias',
-                fallback=False)
-            args.relative_attention = config.getboolean(component,
-                                                        'relative_attention',
-                                                        fallback=True)
-            args.num_buckets = config.getint(component,
-                                             'relative_attention_num_buckets')
-            args.max_distance = config.getint(
-                component, 'relative_attention_max_distance')
-            args.logits_dtype = config.get(component,
-                                           'logits_dtype',
-                                           fallback='float32')
-            args.rescale_before_lm_head = config.getboolean(
-                component, 'tie_word_embeddings'
-            )  # default is True (for T5), but False for Flan-T5
-            args.encoder_hidden_size = config.getint('decoder', 'hidden_size')
-            args.encoder_num_heads = config.getint('decoder', 'num_heads')
-            args.encoder_head_size = config.getint('decoder', 'd_kv')
-            args.position_embedding_type = config.get(
-                'structure', 'position_embedding_type')
-            args.decoder_start_token_id = config.getint(
-                'decoder', 'decoder_start_token_id')
-            args.eos_token_id = config.getint('decoder', 'eos_token_id')
-            bos_token_id = config.get('decoder', 'bos_token_id')
-            # pix2struct does not have bos_token_id
-            args.bos_token_id = int(
-                bos_token_id) if bos_token_id != "None" else None
-            args.pad_token_id = config.getint('decoder', 'pad_token_id')
-
-        else:
-            assert False, 'Unsupported component!'
-        return args
-
-    decoder_args = parse_pix2struct_config_by_component(config, "decoder", args)
-    return None, decoder_args
-
-
-def convert_pix2struct_weights_to_tllm_safetensors(config, component, params):
-    weights = {}
-
-    mapping = config.mapping
-
-    convert_weight_to_dtype(params, config.dtype)
-    hidden_size = config.hidden_size
-    ffn_hidden_size = config.intermediate_size
-    num_layers = config.num_hidden_layers
-    n_head = config.num_attention_heads
-    head_size = config.head_size
-    attention_hidden_size = n_head * head_size  # head size * num_heads not necessarily equals hidden_dim, such as Flan-T5
-
-    hf_param_prefix = f'{component}'
-    trtllm_layer_name = f'{component}_layers'
-    trtllm_attn_layer_name = 'self_attention'
-    trtllm_attn_layernorm_name = 'self_attention_layernorm'
-
-    def get_attn_module_name(component, layer, attn_type):
-        return f'{component}.layer.{int(layer)}.{attn_type}.attention'
-
-    weights['embedding.vocab_embedding.weight'] = reshape(
-        params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), None)
-
-    layers_range = mapping.pp_layers(num_layers)
-    for layer_idx in layers_range:
-        local_layer_idx = layer_idx - layers_range[0]
-        trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}'
-        hf_layer_name_prefix = f'{hf_param_prefix}.layer.{layer_idx}'
-
-        hidden_layer_name_split = {
-            f'{hf_layer_name_prefix}.self_attention.attention.output.weight': {
-                "name":
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.dense.weight',
-                "shape":
-                (hidden_size, attention_hidden_size // mapping.tp_size),
-                "split_dim": -1
-            },
-            f'{hf_layer_name_prefix}.mlp.DenseReluDense.wo.weight': {
-                "name": f'{trtllm_layer_name_prefix}.mlp.proj.weight',
-                "shape": (hidden_size, ffn_hidden_size // mapping.tp_size),
-                "split_dim": -1
-            },
-            f'{hf_layer_name_prefix}.mlp.DenseReluDense.wi_0.weight': {
-                "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight',
-                "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-                "split_dim": 0
-            },
-        }
-
-        hidden_layer_name_no_split = {
-            f'{hf_layer_name_prefix}.self_attention.layer_norm.weight': {
-                "name":
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layernorm_name}.weight',
-                "shape": None
-            },
-            f'{hf_layer_name_prefix}.mlp.layer_norm.weight': {
-                "name": f'{trtllm_layer_name_prefix}.mlp_layernorm.weight',
-                "shape": None
-            },
-        }
-
-        if config.gated_act:
-            hidden_layer_name_split.update({
-                f'{hf_layer_name_prefix}.mlp.DenseReluDense.wi_1.weight': {
-                    "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight',
-                    "shape": (ffn_hidden_size // mapping.tp_size, hidden_size),
-                    "split_dim": 0
-                },
-            })
-
-        hidden_layer_name_split.update({
-            f'{hf_layer_name_prefix}.encoder_decoder_attention.attention.output.weight':
-            {
-                "name":
-                f'{trtllm_layer_name_prefix}.cross_attention.dense.weight',
-                "shape":
-                (hidden_size, attention_hidden_size // mapping.tp_size),
-                "split_dim": -1
-            },
-        })
-        hidden_layer_name_no_split.update({
-            f'{hf_layer_name_prefix}.encoder_decoder_attention.layer_norm.weight':
-            {
-                "name":
-                f'{trtllm_layer_name_prefix}.cross_attention_layernorm.weight',
-                "shape": None
-            },
-        })
-        self_attn_module_name = get_attn_module_name(
-            component, layer_idx, 'encoder_decoder_attention')
-        weights.update(
-            fuse_qkv_one_layer(
-                params, self_attn_module_name,
-                f'{trtllm_layer_name_prefix}.cross_attention', mapping.tp_size,
-                mapping.tp_rank, config.model_type,
-                (attention_hidden_size * 3 // mapping.tp_size, hidden_size),
-                None))
-
-        self_attn_module_name = get_attn_module_name(component, layer_idx,
-                                                     'self_attention')
-        weights.update(
-            fuse_qkv_one_layer(
-                params, self_attn_module_name,
-                f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}',
-                mapping.tp_size, mapping.tp_rank, config.model_type,
-                (attention_hidden_size * 3 // mapping.tp_size, hidden_size),
-                None))
-
-        weights[
-            f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.rel_attn_table'] = reshape(
-                split(
-                    params[
-                        f'{component}.layer.0.self_attention.attention.relative_attention_bias.weight']
-                    .T, mapping.tp_size, mapping.tp_rank, 0),
-                (n_head // mapping.tp_size, config.num_buckets))
-
-        for hf_weight_name, weight_info in hidden_layer_name_split.items():
-            if hf_weight_name in params.keys():
-                weights[weight_info["name"]] = reshape(
-                    split(params[hf_weight_name],
-                          mapping.tp_size,
-                          mapping.tp_rank,
-                          dim=weight_info["split_dim"]), weight_info["shape"])
-        for hf_weight_name, weight_info in hidden_layer_name_no_split.items():
-            if hf_weight_name in params.keys():
-                weights[weight_info["name"]] = reshape(
-                    params[hf_weight_name].clone(), shape=weight_info["shape"])
-
-    weights[f'final_layernorm.weight'] = reshape(
-        params[f'{component}.final_layer_norm.weight'].clone(), None)
-
-    weights['lm_head.weight'] = reshape(
-        split(params[f'{component}.lm_head.weight'],
-              mapping.tp_size,
-              mapping.tp_rank,
-              dim=0), (config.vocab_size // mapping.tp_size, hidden_size))
-    if not config.use_implicit_relative_attention:
-        weights[f'rel_attn_table'] = reshape(
-            split(
-                params[
-                    f'{component}.layer.0.self_attention.attention.relative_attention_bias.weight']
-                .T, mapping.tp_size, mapping.tp_rank, 0),
-            (n_head // mapping.tp_size, config.num_buckets))
-
-    return weights
-
-
-def get_model(args):
-    if args.model_type == "t5":
-        model = T5ForConditionalGeneration.from_pretrained(args.model_dir)
-    elif args.model_type == "nmt":
-        from fairseq.models.transformer import TransformerModel
-        model = TransformerModel.from_pretrained(args.model_dir)
-    elif args.model_type == "bart":
-        if args.nougat:
-            model = VisionEncoderDecoderModel.from_pretrained(args.model_dir)
-            model = model.get_decoder()
-        else:
-            model = AutoModelForSeq2SeqLM.from_pretrained(args.model_dir)
-    elif args.model_type == "pix2struct":
-        model = Pix2StructForConditionalGeneration.from_pretrained(
-            args.model_dir)
-    elif args.model_type == "blip2":
-        model = Blip2ForConditionalGeneration.from_pretrained(
-            args.model_dir).language_model
-    return model
-
-
-def convert_checkpoint(args):
-
-    model = get_model(args)
-
-    saved_dir = Path(args.output_dir)
-    saved_dir.mkdir(parents=True, exist_ok=True)
-
-    encoder_saved_dir = saved_dir / "encoder"
-    encoder_saved_dir.mkdir(parents=True, exist_ok=True)
-    decoder_saved_dir = saved_dir / "decoder"
-    decoder_saved_dir.mkdir(parents=True, exist_ok=True)
-
-    world_size = args.tp_size * args.pp_size
-
-    kv_cache_quant_algo = None
-    quant_algo = None
-
-    model_type = args.model_type if args.model_type != "blip2" else "t5"
-    encoder_config, decoder_config = globals()[f'parse_{model_type}_config'](
-        args, model)
-
-    additional_settings = ["gated_act"]
-
-    if not args.nougat and args.model_type != "pix2struct":
-        tllm_encoder_config = {
-            'architecture': "EncoderModel",
-            'dtype': args.dtype,
-            'logits_dtype': encoder_config.logits_dtype,
-            'num_hidden_layers': encoder_config.n_layer,
-            'num_attention_heads': encoder_config.n_head,
-            'hidden_size': encoder_config.hidden_size,
-            'norm_epsilon': encoder_config.layernorm_eps,
-            'vocab_size': encoder_config.vocab_size,
-            'position_embedding_type': encoder_config.position_embedding_type,
-            'hidden_act': encoder_config.hidden_act,
-            'quantization': {
-                'quant_algo': quant_algo,
-                'kv_cache_quant_algo': kv_cache_quant_algo,
-            },
-            'mapping': {
-                'world_size': world_size,
-                'tp_size': args.tp_size,
-                'pp_size': args.pp_size,
-            },
-            'use_parallel_embedding': args.use_parallel_embedding,
-            'embedding_sharding_dim': args.embedding_sharding_dim,
-            'max_position_embeddings': encoder_config.n_positions,
-            'num_key_value_heads': encoder_config.n_head,
-            'head_size': encoder_config.head_size,
-            'has_position_embedding': encoder_config.has_position_embedding,
-            'layernorm_type': encoder_config.layernorm_type,
-            'has_attention_qkvo_bias': encoder_config.has_attention_qkvo_bias,
-            'has_mlp_bias': encoder_config.has_mlp_bias,
-            'has_model_final_layernorm':
-            encoder_config.has_model_final_layernorm,
-            'has_embedding_layernorm': encoder_config.has_embedding_layernorm,
-            'has_embedding_scale': encoder_config.has_embedding_scale,
-            'intermediate_size': encoder_config.ffn_hidden_size,
-            'q_scaling': encoder_config.q_scaling,
-            'layernorm_position': encoder_config.layernorm_position,
-            'mlp_type': encoder_config.mlp_type,
-            'relative_attention': encoder_config.relative_attention,
-            'max_distance': encoder_config.max_distance,
-            'num_buckets': encoder_config.num_buckets,
-            'model_type': encoder_config.model_type,
-        }
-
-        for additional_setting in additional_settings:
-            if hasattr(encoder_config, additional_setting):
-                tllm_encoder_config.update({
-                    additional_setting:
-                    getattr(encoder_config, additional_setting)
-                })
-
-        with (encoder_saved_dir / "config.json").open('w') as f:
-            json.dump(tllm_encoder_config, f, indent=4)
-
-        encoder_convert_args = dict(params=model.state_dict(),
-                                    component="encoder")
-    tllm_decoder_config = {
-        'architecture': "DecoderModel",
-        'dtype': args.dtype,
-        'logits_dtype': decoder_config.logits_dtype,
-        'num_hidden_layers': decoder_config.n_layer,
-        'num_attention_heads': decoder_config.n_head,
-        'hidden_size': decoder_config.hidden_size,
-        'norm_epsilon': decoder_config.layernorm_eps,
-        'vocab_size': decoder_config.vocab_size,
-        'position_embedding_type': decoder_config.position_embedding_type,
-        'hidden_act': decoder_config.hidden_act,
-        'quantization': {
-            'quant_algo': quant_algo,
-            'kv_cache_quant_algo': kv_cache_quant_algo,
-        },
-        'mapping': {
-            'world_size': world_size,
-            'tp_size': args.tp_size,
-            'pp_size': args.pp_size,
-        },
-        'use_parallel_embedding': args.use_parallel_embedding,
-        'embedding_sharding_dim': args.embedding_sharding_dim,
-        'max_position_embeddings': decoder_config.n_positions,
-        'head_size': decoder_config.head_size,
-        'has_position_embedding': decoder_config.has_position_embedding,
-        'layernorm_type': decoder_config.layernorm_type,
-        'has_attention_qkvo_bias': decoder_config.has_attention_qkvo_bias,
-        'has_mlp_bias': decoder_config.has_mlp_bias,
-        'has_model_final_layernorm': decoder_config.has_model_final_layernorm,
-        'has_embedding_layernorm': decoder_config.has_embedding_layernorm,
-        'has_embedding_scale': decoder_config.has_embedding_scale,
-        'intermediate_size': decoder_config.ffn_hidden_size,
-        'q_scaling': decoder_config.q_scaling,
-        'layernorm_position': decoder_config.layernorm_position,
-        'mlp_type': decoder_config.mlp_type,
-        'relative_attention': decoder_config.relative_attention,
-        'max_distance': decoder_config.max_distance,
-        'num_buckets': decoder_config.num_buckets,
-        'model_type': decoder_config.model_type,
-        'rescale_before_lm_head': decoder_config.rescale_before_lm_head,
-        'encoder_hidden_size': decoder_config.encoder_hidden_size,
-        'encoder_num_heads': decoder_config.encoder_num_heads,
-        'encoder_head_size': decoder_config.encoder_head_size,
-        'skip_cross_kv': args.skip_cross_kv,
-        'use_implicit_relative_attention': args.use_implicit_relative_attention,
-        'decoder_start_token_id': decoder_config.decoder_start_token_id,
-        'eos_token_id': decoder_config.eos_token_id,
-        'bos_token_id': decoder_config.bos_token_id,
-        'pad_token_id': decoder_config.pad_token_id,
-    }
-    for additional_setting in additional_settings:
-        if hasattr(decoder_config, additional_setting):
-            tllm_decoder_config.update({
-                additional_setting:
-                getattr(decoder_config, additional_setting)
-            })
-
-    with (decoder_saved_dir / "config.json").open('w') as f:
-        json.dump(tllm_decoder_config, f, indent=4)
-
-    decoder_convert_args = dict(params=model.state_dict(), component="decoder")
-
-    if args.model_type == "nmt":
-        fairseq_config = vars(model.cfg.model)  # Namespace --> dict
-        num_embeddings = fairseq_config['max_source_positions']
-        embedding_dim = fairseq_config['encoder_embed_dim']
-        padding_idx = model.models[0].encoder.embed_tokens.padding_idx  # 1
-
-        sin_pos_embedding = model.models[
-            0].encoder.embed_positions.get_embedding(
-                padding_idx + 1 + num_embeddings,
-                embedding_dim,
-                padding_idx=padding_idx)  # [2 + num_embeddings, embed_dim]
-        sin_pos_embedding = sin_pos_embedding[2:, :]  # remove offset embeddings
-
-        encoder_convert_args["sin_pos_embedding"] = sin_pos_embedding
-        decoder_convert_args["sin_pos_embedding"] = sin_pos_embedding
-
-    if args.workers == 1:
-        if not args.nougat and args.model_type != "pix2struct":
-            convert(0, world_size, args, tllm_encoder_config,
-                    encoder_convert_args, encoder_saved_dir)
-        convert(0, world_size, args, tllm_decoder_config, decoder_convert_args,
-                decoder_saved_dir)
-    else:
-        if args.workers > world_size:
-            args.workers = world_size
-        LOGGER.info(f'Convert checkpoint using {args.workers} workers.')
-        import torch.multiprocessing as mp
-        if not args.nougat and args.model_type != "pix2struct":
-            mp.spawn(convert,
-                     nprocs=args.workers,
-                     args=(world_size, args, tllm_encoder_config,
-                           encoder_convert_args, encoder_saved_dir))
-        mp.spawn(convert,
-                 nprocs=args.workers,
-                 args=(world_size, args, tllm_decoder_config,
-                       decoder_convert_args, decoder_saved_dir))
-
-
-def convert(worker_rank, world_size, args, model_config, convert_args,
-            saved_dir):
-    for rank in range(worker_rank, world_size, args.workers):
-        rank_config = copy.deepcopy(PretrainedConfig.from_dict(model_config))
-        rank_config.set_rank(rank)
-        weights = globals(
-        )[f'convert_{rank_config.model_type}_weights_to_tllm_safetensors'](
-            config=rank_config, **convert_args)
-        safetensors.torch.save_file(weights,
-                                    f'{saved_dir}/rank{rank}.safetensors')
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.RawTextHelpFormatter)
-    parser.add_argument(
-        '--model_type',
-        type=str,
-        default='t5',
-        choices=['t5', 'nmt', 'bart', 'pix2struct', 'blip2'],
-        help=
-        'Multimodal type when this script is used for multimodal conversion.')
-
-    parser.add_argument('--tp_size',
-                        type=int,
-                        default=1,
-                        help='N-way tensor parallelism size')
-    parser.add_argument('--pp_size',
-                        type=int,
-                        default=1,
-                        help='N-way pipeline parallelism size')
-    parser.add_argument("--model_dir",
-                        "-i",
-                        type=str,
-                        help="Path to the framework checkpoint file",
-                        required=True)
-    parser.add_argument("--output_dir",
-                        "-o",
-                        type=str,
-                        help="Path to the converted TRT-LLM model weight file",
-                        required=True)
-    parser.add_argument(
-        "--workers",
-        type=int,
-        help="How many workers to spawn for conversion (default: 4)",
-        default=4)
-    parser.add_argument("--nougat",
-                        action="store_true",
-                        help="Model which uses vision encoder + mbart decoder")
-    parser.add_argument("--verbose",
-                        action="store_true",
-                        help="Provide verbose messages")
-    parser.add_argument(
-        '--use_parallel_embedding',
-        action="store_true",
-        default=False,
-        help=
-        'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled'
-    )
-    parser.add_argument(
-        '--embedding_sharding_dim',
-        type=int,
-        default=0,
-        choices=[0, 1],
-        help=
-        'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). '
-        'To shard it along hidden dimension, set embedding_sharding_dim=1'
-        'Note: embedding sharding is only enabled when embedding_sharding_dim = 0'
-    )
-    parser.add_argument(
-        '--use_weight_only',
-        default=False,
-        action="store_true",
-        help='Quantize weights for the various GEMMs to INT4/INT8.'
-        'See --weight_only_precision to set the precision')
-    parser.add_argument(
-        '--weight_only_precision',
-        const='int8',
-        type=str,
-        nargs='?',
-        default='int8',
-        choices=['int8', 'int4'],
-        help=
-        'Define the precision for the weights when using weight-only quantization.'
-        'You must also use --use_weight_only for that argument to have an impact.'
-    )
-    parser.add_argument(
-        '--dtype',
-        type=str,
-        default='float16',
-        choices=['float16', 'float32', 'bfloat16'],
-        help=
-        'Target inference dtype. Weights and Computation will be in this dtype, no matter what original dtype the weight checkpoint has.'
-    )
-    parser.add_argument(
-        '--skip_cross_kv',
-        action='store_true',
-        help=
-        'Skip redundant cross qkv computation by using TensorRT IfConditional switch (experimental).'
-    )
-    parser.add_argument(
-        '--use_implicit_relative_attention',
-        action='store_true',
-        help=
-        'Compute relative attention bias on the fly instead of pre-compute a relative attention bias table.'
-    )
-    args = parser.parse_args()
-    log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s"
-    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO,
-                        format=log_format)
-    LOGGER.info("\n=============== Argument ===============")
-    for key in vars(args):
-        LOGGER.info(f"{key}: {vars(args)[key]}")
-    LOGGER.info("========================================")
-
-    start_time = datetime.now()
-    convert_checkpoint(args)
-    stop_time = datetime.now()
-    run_time = (stop_time - start_time)
-    LOGGER.info("Spend {} (h:m:s) to convert the model".format(run_time))

deployment/tensorrt_llm/convert/helper.py

@@ -1,95 +0,0 @@
-# copied from: https://github.com/NVIDIA/TensorRT-LLM/blob/v0.18.1/examples/enc_dec/helper.py
-
-import typing
-from typing import Union
-
-import numpy as np
-import torch  # pytype: disable=import-error
-
-from tensorrt_llm._utils import str_dtype_to_torch
-
-
-def split(v: Union[np.ndarray, torch.Tensor],
-          tp_size: int,
-          tp_rank: int,
-          dim=0):
-    if tp_size == 1:
-        if isinstance(v, np.ndarray):
-            return np.ascontiguousarray(v.copy())
-        else:
-            return v.clone().detach()
-    assert len(v.shape) > 1 or dim == 0
-    if isinstance(v, np.ndarray):
-        return np.ascontiguousarray(
-            np.split(v, tp_size, axis=dim)[tp_rank].copy())
-    else:
-        assert v.shape[dim] % tp_size == 0, \
-            'Unable to split: shape={v.shape} (dim={dim}) tp_size={tp_size}.'
-        split_size = v.shape[dim] // tp_size
-        return v.split(split_size, dim=dim)[tp_rank].clone().detach()
-
-
-def reshape(v: torch.Tensor, shape=None):
-    if shape is None:
-        return v.contiguous()
-    else:
-        return v.reshape(shape).contiguous()
-
-
-def fuse_qkv_one_layer(params, attn_module_name, trtllm_layer_name, tp_size,
-                       tp_rank, model_type, weight_shape, bias_shape):
-
-    qkv_module_names = get_qkv_module_name(model_type)
-
-    weight = {}
-
-    # fuse weights of q, k, v
-    q_w = params[f'{attn_module_name}.{qkv_module_names["q"]}.weight']
-    k_w = params[f'{attn_module_name}.{qkv_module_names["k"]}.weight']
-    v_w = params[f'{attn_module_name}.{qkv_module_names["v"]}.weight']
-
-    # fuse qkv weight
-    shape = q_w.shape  # (do, din)
-    qkv_w = torch.cat([q_w, k_w, v_w],
-                      dim=0).reshape([3, shape[0], shape[1]])  # (3, do, din)
-    qkv_w = split(qkv_w, tp_size, tp_rank, dim=1)
-    weight[f'{trtllm_layer_name}.qkv.weight'] = reshape(qkv_w,
-                                                        shape=weight_shape)
-
-    # fuse qkv biases if present
-    if f'{attn_module_name}.{qkv_module_names["q"]}.bias' in params.keys(
-    ) and params[f'{attn_module_name}.{qkv_module_names["q"]}.bias'] is not None:
-        q_b = params[f'{attn_module_name}.{qkv_module_names["q"]}.bias']
-        k_b = params[f'{attn_module_name}.{qkv_module_names["k"]}.bias']
-        v_b = params[f'{attn_module_name}.{qkv_module_names["v"]}.bias']
-        shape = q_b.shape[0]  # (do,)
-        qkv_b = torch.cat([q_b, k_b, v_b], dim=0).reshape([3, shape])  # (3, do)
-        qkv_b = split(qkv_b, tp_size, tp_rank, dim=1)
-        weight[f'{trtllm_layer_name}.qkv.bias'] = reshape(qkv_b,
-                                                          shape=bias_shape)
-    return weight
-
-
-def get_qkv_module_name(model_type):
-    if model_type in ["t5", "blip2"]:
-        q = "q"
-        k = "k"
-        v = "v"
-    elif model_type == "bart" or model_type == "nmt":
-        q = "q_proj"
-        k = "k_proj"
-        v = "v_proj"
-    elif model_type == "pix2struct":
-        q = "query"
-        k = "key"
-        v = "value"
-    return {"q": q, "k": k, "v": v}
-
-
-def convert_weight_to_dtype(params: typing.Dict[str, torch.Tensor],
-                            dtype: typing.Optional[np.dtype] = None):
-    if dtype is not None:
-        assert isinstance(dtype,
-                          str), f"dtype must be str, but get type {type(dtype)}"
-        for name in params.keys():
-            params[name] = params[name].to(str_dtype_to_torch(dtype))

deployment/tensorrt_llm/convert_dolphin.sh

@@ -1,47 +0,0 @@
-#!/usr/bin/env bash
-set -ex
-
-############################################################################################
-# Reference: https://github.com/NVIDIA/TensorRT-LLM/tree/v0.18.2/examples/multimodal#nougat
-############################################################################################
-
-export LD_LIBRARY_PATH=/usr/local/lib/python3.10/site-packages/tensorrt_libs/:/usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib/:$LD_LIBRARY_PATH
-
-# 1. Download Huggingface weights
-export MODEL_NAME="Dolphin"
-git clone https://huggingface.co/Bytedance/${MODEL_NAME} tmp/hf_models/${MODEL_NAME}
-
-
-export MAX_BATCH_SIZE=16
-export MAX_SEQ_LEN=4096
-export MAX_INPUT_LEN=10
-export MAX_ENCODER_INPUT_LEN=784
-
-# 2. Convert Huggingface weights into TRT-LLM checkpoints and build TRT engines using scripts in examples/enc_dec
-python ./convert/convert_checkpoint.py --model_type bart \
-    --model_dir tmp/hf_models/${MODEL_NAME} \
-    --output_dir tmp/trt_models/${MODEL_NAME}/bfloat16 \
-    --tp_size 1 \
-    --pp_size 1 \
-    --dtype bfloat16 \
-    --nougat
-
-
-trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/bfloat16/decoder \
-    --output_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16/decoder \
-    --paged_kv_cache disable \
-    --moe_plugin disable \
-    --gemm_plugin bfloat16 \
-    --bert_attention_plugin bfloat16 \
-    --gpt_attention_plugin bfloat16 \
-    --remove_input_padding enable \
-    --max_beam_width 1 \
-    --max_batch_size ${MAX_BATCH_SIZE} \
-    --max_seq_len ${MAX_SEQ_LEN} \
-    --max_input_len ${MAX_INPUT_LEN} \
-    --max_encoder_input_len $((${MAX_BATCH_SIZE} * ${MAX_ENCODER_INPUT_LEN})) # MAX_BATCH_SIZE (max_batch_size) * MAX_ENCODER_INPUT_LEN (num_visual_features)
-
-# 3. Generate TensorRT engines for visual components and combine everything into final pipeline.
-python ./convert/build_visual_engine.py --model_type nougat \
-    --model_path tmp/hf_models/${MODEL_NAME} \
-    --max_batch_size ${MAX_BATCH_SIZE}

deployment/tensorrt_llm/dolphin_runner.py

@@ -1,220 +0,0 @@
-"""
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
-import json
-import os
-from typing import Optional
-
-import tensorrt_llm
-import tensorrt_llm.profiler as profiler
-import torch
-from PIL import Image
-from pydantic import BaseModel, Field
-from tensorrt_llm import logger
-from tensorrt_llm import mpi_rank
-from tensorrt_llm.runtime import MultimodalModelRunner
-from transformers import AutoTokenizer, DonutProcessor
-
-
-class InferenceConfig(BaseModel):
-    max_new_tokens: int = Field(128, description="Maximum new tokens to generate")
-    batch_size: int = Field(1, description="Batch size for inference")
-    log_level: str = Field("info", description="Logging level")
-    visual_engine_dir: Optional[str] = Field(None, description="Directory for visual engine files")
-    visual_engine_name: str = Field("model.engine", description="Visual engine filename")
-    llm_engine_dir: Optional[str] = Field(None, description="Directory for LLM engine files")
-    hf_model_dir: Optional[str] = Field(None, description="Hugging Face model directory")
-    input_text: Optional[str] = Field(None, description="Input text for inference")
-    num_beams: int = Field(1, description="Number of beams for beam search")
-    top_k: int = Field(1, description="Top-k sampling value")
-    top_p: float = Field(0.0, description="Top-p (nucleus) sampling value")
-    temperature: float = Field(1.0, description="Sampling temperature")
-    repetition_penalty: float = Field(1.0, description="Repetition penalty factor")
-    run_profiling: bool = Field(False, description="Enable profiling mode")
-    profiling_iterations: int = Field(20, description="Number of profiling iterations")
-    check_accuracy: bool = Field(False, description="Enable accuracy checking")
-    video_path: Optional[str] = Field(None, description="Path to input video file")
-    video_num_frames: Optional[int] = Field(None, description="Number of video frames to process")
-    image_path: Optional[str] = Field(None, description="Path to input image file")
-    path_sep: str = Field(",", description="Path separator character")
-    prompt_sep: str = Field(",", description="Prompt separator character")
-    enable_context_fmha_fp32_acc: Optional[bool] = Field(
-        None,
-        description="Enable FP32 accumulation for context FMHA"
-    )
-    enable_chunked_context: bool = Field(False, description="Enable chunked context processing")
-    use_py_session: bool = Field(False, description="Use Python session instead of C++")
-    kv_cache_free_gpu_memory_fraction: float = Field(
-        0.9,
-        description="Fraction of GPU memory free for KV cache",
-        ge=0.0, le=1.0
-    )
-    cross_kv_cache_fraction: float = Field(
-        0.5,
-        description="Fraction of cross-attention KV cache",
-        ge=0.0, le=1.0
-    )
-    multi_block_mode: bool = Field(True, description="Enable multi-block processing mode")
-
-
-class DolphinRunner(MultimodalModelRunner):
-    def __init__(self, args):
-        self.args = args
-
-        self.runtime_rank = mpi_rank()
-        device_id = self.runtime_rank % torch.cuda.device_count()
-        torch.cuda.set_device(device_id)
-        self.device = "cuda:%d" % (device_id)
-
-        self.stream = torch.cuda.Stream(torch.cuda.current_device())
-        torch.cuda.set_stream(self.stream)
-
-        # parse model type from visual engine config
-        with open(os.path.join(self.args.visual_engine_dir, "config.json"),
-                  "r") as f:
-            config = json.load(f)
-        self.model_type = config['builder_config']['model_type']
-        self.vision_precision = config['builder_config']['precision']
-        self.decoder_llm = not (
-                't5' in self.model_type
-                or self.model_type in ['nougat', 'pix2struct']
-        )  # BLIP2-T5, pix2struct and Nougat are using encoder-decoder models as LLMs
-
-        if self.model_type == "mllama":
-            self.vision_input_names = [
-                "pixel_values",
-                "aspect_ratio_ids",
-                "aspect_ratio_mask",
-            ]
-            self.vision_output_names = [
-                "output",
-            ]
-        else:
-            self.vision_input_names = ["input"]
-            self.vision_output_names = ["output"]
-
-        self.use_py_session = True
-
-        self.init_image_encoder()
-        self.init_tokenizer()
-        self.init_processor()
-        self.init_llm()
-
-    def init_tokenizer(self):
-        assert self.model_type == 'nougat'
-        self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_model_dir)
-        self.tokenizer.padding_side = "right"
-
-    def init_processor(self):
-        assert self.model_type == 'nougat'
-        self.processor = DonutProcessor.from_pretrained(self.args.hf_model_dir, use_fast=True)
-
-    def run(self, input_texts, input_images, max_new_tokens):
-        prompts = [f"<s>{text.strip()} <Answer/>" for text in input_texts]
-        images = self.processor(input_images, return_tensors="pt")['pixel_values'].to("cuda")
-        prompt_ids = self.tokenizer(prompts, add_special_tokens=False, return_tensors="pt").input_ids.to("cuda")
-
-        # 🚨🚨🚨 Important! If the type of prompt_ids is not int32, the output will be wrong. 🚨🚨🚨
-        prompt_ids = prompt_ids.to(torch.int32)
-
-        logger.info("---------------------------------------------------------")
-        logger.info(f"images size: {images.size()}")
-        logger.info(f"prompt_ids: {prompt_ids}, size: {prompt_ids.size()}, dtype: {prompt_ids.dtype}")
-        logger.info("---------------------------------------------------------")
-
-        output_texts = self.generate(input_texts,
-                                     [None] * len(input_texts),
-                                     images,
-                                     prompt_ids,
-                                     max_new_tokens,
-                                     warmup=False,
-                                     )
-
-        return output_texts
-
-    def generate(self,
-                 pre_prompt,
-                 post_prompt,
-                 image,
-                 decoder_input_ids,
-                 max_new_tokens,
-                 warmup=False,
-                 other_vision_inputs={},
-                 other_decoder_inputs={}):
-        if not warmup:
-            profiler.start("Generate")
-        input_ids, input_lengths, ptuning_args, visual_features = self.preprocess(
-            warmup, pre_prompt, post_prompt, image, other_vision_inputs)
-
-        if warmup: return None
-
-        # use prompt tuning to pass multimodal features
-        # model.generate() expects the following params (see layers/embedding.py):
-        # args[0]: prompt embedding table, [batch_size, multimodal_len, hidden_size], later flattened to [batch_size * multimodal_len, hidden_size]
-        # args[1]: prompt task ids, [batch_size]. in multimodal case, arange(batch_size), i.e. in VILA batching mode 2, each image is treated separately in the batch instead of concated together (although the prompt embedding table has to be concated)
-        # args[2]: prompt task vocab size, [1]. assuming all table has the same length, which in multimodal case equals to multimodal_len
-        profiler.start("LLM")
-        if self.model_type in ['nougat', 'pix2struct']:
-            # Trim encoder input_ids to match visual features shape
-            ids_shape = (min(self.args.batch_size, len(pre_prompt)), visual_features.shape[1])
-            if self.model_type == 'nougat':
-                input_ids = torch.zeros(ids_shape, dtype=torch.int32)
-            elif self.model_type == 'pix2struct':
-                input_ids = torch.ones(ids_shape, dtype=torch.int32)
-
-        output_ids = self.model.generate(
-            input_ids,
-            decoder_input_ids,
-            max_new_tokens,
-            num_beams=self.args.num_beams,
-            bos_token_id=self.tokenizer.bos_token_id,
-            pad_token_id=self.tokenizer.pad_token_id,
-            eos_token_id=self.tokenizer.eos_token_id,
-            debug_mode=False,
-            prompt_embedding_table=ptuning_args[0],
-            prompt_tasks=ptuning_args[1],
-            prompt_vocab_size=ptuning_args[2],
-        )
-        profiler.stop("LLM")
-
-        if mpi_rank() == 0:
-            # Extract a list of tensors of shape beam_width x output_ids.
-            output_beams_list = [
-                self.tokenizer.batch_decode(
-                    output_ids[batch_idx, :, decoder_input_ids.shape[1]:],
-                    skip_special_tokens=False) for batch_idx in range(
-                    min(self.args.batch_size, decoder_input_ids.shape[0]))
-            ]
-
-            stripped_text = [[
-                output_beams_list[batch_idx][beam_idx].replace("</s>", "").replace("<pad>", "").strip()
-                for beam_idx in range(self.args.num_beams)
-            ] for batch_idx in range(
-                min(self.args.batch_size, decoder_input_ids.shape[0]))]
-            profiler.stop("Generate")
-            return stripped_text
-        else:
-            profiler.stop("Generate")
-            return None
-
-
-if __name__ == "__main__":
-    config = InferenceConfig(
-        max_new_tokens=4024,
-        batch_size=16,
-        log_level="info",
-        hf_model_dir=f"./tmp/hf_models/Dolphin",
-        visual_engine_dir=f"./tmp/trt_engines/Dolphin/vision_encoder",
-        llm_engine_dir=f"./tmp/trt_engines/Dolphin/1-gpu/bfloat16",
-    )
-
-    model = DolphinRunner(config)
-
-    image_path = "../../demo/page_imgs/page_1.jpeg"
-    prompt = "Parse the reading order of this document."
-    image = Image.open(image_path).convert("RGB")
-    output_texts = model.run([prompt], [image], 4024)
-    output_texts = [texts[0] for texts in output_texts]
-    print(output_texts)