local

.gitattributes

@@ -33,8 +33,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-*.png filter=lfs diff=lfs merge=lfs -text
-*.jpeg filter=lfs diff=lfs merge=lfs -text
-*.jpg filter=lfs diff=lfs merge=lfs -text
-*.wav filter=lfs diff=lfs merge=lfs -text
-*.gif filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,1084 +0,0 @@
-import gradio as gr
-import json
-import markdown
-import cv2
-import numpy as np
-from PIL import Image
-from transformers import AutoProcessor, VisionEncoderDecoderModel, AutoModelForImageTextToText
-import torch
-try:
-    from sentence_transformers import SentenceTransformer
-    import numpy as np
-    from sklearn.metrics.pairwise import cosine_similarity
-    RAG_DEPENDENCIES_AVAILABLE = True
-except ImportError as e:
-    print(f"RAG dependencies not available: {e}")
-    print("Please install: pip install sentence-transformers scikit-learn")
-    RAG_DEPENDENCIES_AVAILABLE = False
-    SentenceTransformer = None
-import os
-import tempfile
-import uuid
-import base64
-import io
-from utils.utils import *
-from utils.markdown_utils import MarkdownConverter
-
-# Voice functionality imports
-import time
-import librosa
-from dataclasses import dataclass, field
-from pydub import AudioSegment
-try:
-    from voice_chat.utils.vad import get_speech_timestamps, collect_chunks, VadOptions
-    from voice_chat.gemma3n_inference import Gemma3nInference
-    VOICE_DEPENDENCIES_AVAILABLE = True
-except ImportError as e:
-    print(f"Voice dependencies not available: {e}")
-    VOICE_DEPENDENCIES_AVAILABLE = False
-
-# Math extension is optional for enhanced math rendering
-MATH_EXTENSION_AVAILABLE = False
-try:
-    from mdx_math import MathExtension
-    MATH_EXTENSION_AVAILABLE = True
-except ImportError:
-    pass
-
-# 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 loaded and warmed up successfully")
-    except Exception as e:
-        print(f"⚠️ Voice model initialization failed: {e}")
-        voice_model = None
-
-
-class DOLPHIN:
-    def __init__(self, model_id_or_path):
-        """Initialize the Hugging Face model optimized for powerful GPU"""
-        self.processor = AutoProcessor.from_pretrained(model_id_or_path)
-        self.model = VisionEncoderDecoderModel.from_pretrained(
-            model_id_or_path,
-            torch_dtype=torch.float16,
-            device_map="auto" if torch.cuda.is_available() else None
-        )
-        self.model.eval()
-        
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        if not torch.cuda.is_available():
-            self.model = self.model.float()
-        
-        self.tokenizer = self.processor.tokenizer
-        
-    def chat(self, prompt, image):
-        """Process an image or batch of images with the given prompt(s)"""
-        is_batch = isinstance(image, list)
-        
-        if not is_batch:
-            images = [image]
-            prompts = [prompt]
-        else:
-            images = image
-            prompts = prompt if isinstance(prompt, list) else [prompt] * len(images)
-        
-        batch_inputs = self.processor(images, return_tensors="pt", padding=True)
-        batch_pixel_values = batch_inputs.pixel_values
-        
-        if torch.cuda.is_available():
-            batch_pixel_values = batch_pixel_values.half().to(self.device)
-        else:
-            batch_pixel_values = batch_pixel_values.to(self.device)
-        
-        prompts = [f"<s>{p} <Answer/>" for p in prompts]
-        batch_prompt_inputs = self.tokenizer(
-            prompts,
-            add_special_tokens=False,
-            return_tensors="pt"
-        )
-
-        batch_prompt_ids = batch_prompt_inputs.input_ids.to(self.device)
-        batch_attention_mask = batch_prompt_inputs.attention_mask.to(self.device)
-        
-        with torch.no_grad():
-            outputs = self.model.generate(
-                pixel_values=batch_pixel_values,
-                decoder_input_ids=batch_prompt_ids,
-                decoder_attention_mask=batch_attention_mask,
-                min_length=1,
-                max_length=2048,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id,
-                use_cache=True,
-                bad_words_ids=[[self.tokenizer.unk_token_id]],
-                return_dict_in_generate=True,
-                do_sample=False,
-                num_beams=1,
-                repetition_penalty=1.1,
-                temperature=1.0
-            )
-        
-        sequences = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
-        
-        results = []
-        for i, sequence in enumerate(sequences):
-            cleaned = sequence.replace(prompts[i], "").replace("<pad>", "").replace("</s>", "").strip()
-            results.append(cleaned)
-            
-        if not is_batch:
-            return results[0]
-        return results
-
-
-class Gemma3nModel:
-    def __init__(self, model_id="google/gemma-3n-E4B-it"):
-        """Initialize the Gemma 3n model for text generation and image description"""
-        self.model_id = model_id
-        self.processor = AutoProcessor.from_pretrained(model_id)
-        self.model = AutoModelForImageTextToText.from_pretrained(
-            model_id,
-            torch_dtype="auto",
-            device_map="auto"
-        )
-        self.model.eval()
-        print(f"✅ Gemma 3n loaded (Device: {self.model.device}, DType: {self.model.dtype})")
-        
-    def generate_alt_text(self, pil_image):
-        """Generate alt text for an image using local Gemma 3n"""
-        try:
-            # Ensure image is in RGB mode
-            if pil_image.mode != 'RGB':
-                pil_image = pil_image.convert('RGB')
-            
-            # Create a detailed prompt for alt text generation
-            prompt = """You are an accessibility expert creating alt text for images to help visually impaired users understand visual content. Analyze this image and provide a clear, concise description that captures the essential visual information.
-
-Focus on:
-- Main subject or content of the image
-- Important details, text, or data shown
-- Layout and structure if relevant (charts, diagrams, tables)
-- Context that would help someone understand the image's purpose
-
-Provide a descriptive alt text in 1-2 sentences that is informative but not overly verbose. Start directly with the description without saying "This image shows" or similar phrases."""
-            
-            # Prepare the message format
-            message = {
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": pil_image},
-                    {"type": "text", "text": prompt}
-                ]
-            }
-            
-            # Apply chat template and generate
-            input_ids = self.processor.apply_chat_template(
-                [message],
-                add_generation_prompt=True,
-                tokenize=True,
-                return_dict=True,
-                return_tensors="pt",
-            )
-            input_len = input_ids["input_ids"].shape[-1]
-            
-            input_ids = input_ids.to(self.model.device, dtype=self.model.dtype)
-            outputs = self.model.generate(
-                **input_ids,
-                max_new_tokens=256,
-                disable_compile=True,
-                do_sample=False,
-                temperature=0.1
-            )
-            
-            text = self.processor.batch_decode(
-                outputs[:, input_len:],
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=True
-            )
-            
-            alt_text = text[0].strip()
-            
-            # Clean up the alt text
-            alt_text = alt_text.replace('\n', ' ').replace('\r', ' ')
-            # Remove common prefixes if they appear
-            prefixes_to_remove = ["This image shows", "The image shows", "This shows", "The figure shows"]
-            for prefix in prefixes_to_remove:
-                if alt_text.startswith(prefix):
-                    alt_text = alt_text[len(prefix):].strip()
-                    break
-            
-            return alt_text if alt_text else "Image description unavailable"
-            
-        except Exception as e:
-            print(f"❌ Error generating alt text: {e}")
-            import traceback
-            traceback.print_exc()
-            return "Image description unavailable"
-    
-    def chat(self, prompt, history=None):
-        """Chat functionality using Gemma 3n for text-only conversations"""
-        try:
-            # Create message format
-            message = {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt}
-                ]
-            }
-            
-            # If history exists, include it
-            conversation = history if history else []
-            conversation.append(message)
-            
-            # Apply chat template and generate
-            input_ids = self.processor.apply_chat_template(
-                conversation,
-                add_generation_prompt=True,
-                tokenize=True,
-                return_dict=True,
-                return_tensors="pt",
-            )
-            input_len = input_ids["input_ids"].shape[-1]
-            
-            input_ids = input_ids.to(self.model.device, dtype=self.model.dtype)
-            outputs = self.model.generate(
-                **input_ids,
-                max_new_tokens=1024,
-                disable_compile=True,
-                do_sample=False,
-                pad_token_id=self.processor.tokenizer.pad_token_id
-            )
-            
-            text = self.processor.batch_decode(
-                outputs[:, input_len:],
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=True
-            )
-            
-            return text[0].strip()
-            
-        except Exception as e:
-            print(f"❌ Error in chat: {e}")
-            import traceback
-            traceback.print_exc()
-            return f"Error generating response: {str(e)}"
-
-
-def convert_pdf_to_images_gradio(pdf_file):
-    """Convert uploaded PDF file to list of PIL Images"""
-    try:
-        import pymupdf
-        
-        if isinstance(pdf_file, str):
-            pdf_document = pymupdf.open(pdf_file)
-        else:
-            pdf_bytes = pdf_file.read()
-            pdf_document = pymupdf.open(stream=pdf_bytes, filetype="pdf")
-        
-        images = []
-        for page_num in range(len(pdf_document)):
-            page = pdf_document[page_num]
-            mat = pymupdf.Matrix(2.0, 2.0)
-            pix = page.get_pixmap(matrix=mat)
-            img_data = pix.tobytes("png")
-            pil_image = Image.open(io.BytesIO(img_data)).convert("RGB")
-            images.append(pil_image)
-        
-        pdf_document.close()
-        return images
-        
-    except Exception as e:
-        raise Exception(f"Error converting PDF: {str(e)}")
-
-
-def process_pdf_document(pdf_file, model, progress=gr.Progress()):
-    """Process uploaded PDF file page by page"""
-    if pdf_file is None:
-        return "No PDF file uploaded", ""
-    
-    try:
-        progress(0.1, desc="Converting PDF to images...")
-        images = convert_pdf_to_images_gradio(pdf_file)
-        
-        if not images:
-            return "Failed to convert PDF to images", ""
-        
-        all_results = []
-        
-        for page_idx, pil_image in enumerate(images):
-            progress((page_idx + 1) / len(images) * 0.8 + 0.1, 
-                    desc=f"Processing page {page_idx + 1}/{len(images)}...")
-            
-            layout_output = model.chat("Parse the reading order of this document.", pil_image)
-            
-            padded_image, dims = prepare_image(pil_image)
-            recognition_results = process_elements_optimized(
-                layout_output, 
-                padded_image, 
-                dims, 
-                model, 
-                max_batch_size=4
-            )
-            
-            try:
-                markdown_converter = MarkdownConverter()
-                markdown_content = markdown_converter.convert(recognition_results)
-            except:
-                markdown_content = generate_fallback_markdown(recognition_results)
-            
-            page_result = {
-                "page_number": page_idx + 1,
-                "markdown": markdown_content
-            }
-            all_results.append(page_result)
-        
-        progress(1.0, desc="Processing complete!")
-        
-        combined_markdown = "\n\n---\n\n".join([
-            f"# Page {result['page_number']}\n\n{result['markdown']}" 
-            for result in all_results
-        ])
-        
-        return combined_markdown, "processing_complete"
-        
-    except Exception as e:
-        error_msg = f"Error processing PDF: {str(e)}"
-        return error_msg, "error"
-
-
-def process_elements_optimized(layout_results, padded_image, dims, model, max_batch_size=4):
-    """Optimized element processing for powerful GPU"""
-    layout_results = parse_layout_string(layout_results)
-    
-    text_elements = []
-    table_elements = []
-    figure_results = []
-    previous_box = None
-    reading_order = 0
-    
-    for bbox, label in layout_results:
-        try:
-            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
-                bbox, padded_image, dims, previous_box
-            )
-            
-            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))
-                    pil_crop = crop_margin(pil_crop)
-                    
-                    # Generate alt text for accessibility using local Gemma 3n
-                    alt_text = gemma_model.generate_alt_text(pil_crop)
-                    
-                    buffered = io.BytesIO()
-                    pil_crop.save(buffered, format="PNG")
-                    img_base64 = base64.b64encode(buffered.getvalue()).decode()
-                    data_uri = f"data:image/png;base64,{img_base64}"
-                    
-                    figure_results.append({
-                        "label": label,
-                        "text": f"![{alt_text}]({data_uri})\n\n*{alt_text}*",
-                        "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                        "reading_order": reading_order,
-                        "alt_text": alt_text,
-                    })
-                else:
-                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-                    element_info = {
-                        "crop": pil_crop,
-                        "label": label,
-                        "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                        "reading_order": reading_order,
-                    }
-                    
-                    if label == "tab":
-                        table_elements.append(element_info)
-                    else:
-                        text_elements.append(element_info)
-                        
-            reading_order += 1
-            
-        except Exception as e:
-            print(f"Error processing element {label}: {str(e)}")
-            continue
-    
-    recognition_results = figure_results.copy()
-    
-    if text_elements:
-        text_results = process_element_batch_optimized(
-            text_elements, model, "Read text in the image.", max_batch_size
-        )
-        recognition_results.extend(text_results)
-    
-    if table_elements:
-        table_results = process_element_batch_optimized(
-            table_elements, model, "Parse the table in the image.", max_batch_size
-        )
-        recognition_results.extend(table_results)
-    
-    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
-    return recognition_results
-
-
-def process_element_batch_optimized(elements, model, prompt, max_batch_size=4):
-    """Process elements in batches for powerful GPU"""
-    results = []
-    batch_size = min(len(elements), max_batch_size)
-    
-    for i in range(0, len(elements), batch_size):
-        batch_elements = elements[i:i+batch_size]
-        crops_list = [elem["crop"] for elem in batch_elements]
-        prompts_list = [prompt] * len(crops_list)
-        
-        batch_results = model.chat(prompts_list, crops_list)
-        
-        for j, result in enumerate(batch_results):
-            elem = batch_elements[j]
-            results.append({
-                "label": elem["label"],
-                "bbox": elem["bbox"],
-                "text": result.strip(),
-                "reading_order": elem["reading_order"],
-            })
-            
-        del crops_list, batch_elements
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-    
-    return results
-
-
-def generate_fallback_markdown(recognition_results):
-    """Generate basic markdown if converter fails"""
-    markdown_content = ""
-    for element in recognition_results:
-        if element["label"] == "tab":
-            markdown_content += f"\n\n{element['text']}\n\n"
-        elif element["label"] in ["para", "title", "sec", "sub_sec"]:
-            markdown_content += f"{element['text']}\n\n"
-        elif element["label"] == "fig":
-            # Image should already have alt text from processing
-            markdown_content += f"{element['text']}\n\n"
-    return markdown_content
-
-
-# Initialize models
-model_path = "./hf_model"
-if not os.path.exists(model_path):
-    model_path = "ByteDance/DOLPHIN"
-
-# Model paths and configuration
-model_path = "./hf_model" if os.path.exists("./hf_model") else "ByteDance/DOLPHIN"
-hf_token = os.getenv('HF_TOKEN')
-gemma_model_id = "google/gemma-3n-E4B-it"
-
-# Initialize models
-print("Loading DOLPHIN model...")
-dolphin_model = DOLPHIN(model_path)
-print(f"✅ DOLPHIN model loaded (Device: {dolphin_model.device})")
-
-print("Loading Gemma 3n model...")
-gemma_model = Gemma3nModel(gemma_model_id)
-
-model_status = "✅ Both models loaded successfully"
-
-# Initialize embedding model
-if RAG_DEPENDENCIES_AVAILABLE:
-    try:
-        print("Loading embedding model for RAG...")
-        embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
-        print("✅ Embedding model loaded successfully (CPU)")
-    except Exception as e:
-        print(f"❌ Error loading embedding model: {e}")
-        embedding_model = None
-else:
-    print("❌ RAG dependencies not available")
-    embedding_model = None
-
-
-# Global state for managing tabs
-processed_markdown = ""
-show_results_tab = False
-document_chunks = []
-document_embeddings = None
-
-# Voice chat parameters and state
-IN_CHANNELS = 1
-IN_RATE = 24000
-IN_CHUNK = 1024
-IN_SAMPLE_WIDTH = 2
-VAD_STRIDE = 0.5
-OUT_CHANNELS = 1
-OUT_RATE = 24000
-OUT_SAMPLE_WIDTH = 2
-OUT_CHUNK = 20 * 4096
-
-# Voice model already initialized earlier in the file
-
-@dataclass
-class VoiceAppState:
-    stream: np.ndarray | None = None
-    sampling_rate: int = 0
-    pause_detected: bool = False
-    started_talking: bool = False
-    stopped: bool = False
-    conversation: list = field(default_factory=list)
-
-
-# Voice functionality
-def run_vad(ori_audio, sr):
-    """Voice Activity Detection"""
-    _st = time.time()
-    try:
-        audio = ori_audio
-        if isinstance(audio, bytes):
-            audio = np.frombuffer(audio, dtype=np.int16)
-        audio = audio.astype(np.float32) / 32768.0
-        sampling_rate = 16000
-        if sr != sampling_rate:
-            audio = librosa.resample(audio, orig_sr=sr, target_sr=sampling_rate)
-
-        vad_parameters = {}
-        vad_parameters = VadOptions(**vad_parameters)
-        speech_chunks = get_speech_timestamps(audio, vad_parameters)
-        audio = collect_chunks(audio, speech_chunks)
-        duration_after_vad = audio.shape[0] / sampling_rate
-
-        if sr != sampling_rate:
-            vad_audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=sr)
-        else:
-            vad_audio = audio
-        vad_audio = np.round(vad_audio * 32768.0).astype(np.int16)
-        vad_audio_bytes = vad_audio.tobytes()
-
-        return duration_after_vad, vad_audio_bytes, round(time.time() - _st, 4)
-    except Exception as e:
-        msg = f"[asr vad error] audio_len: {len(ori_audio)/(sr*2):.3f} s, trace: {e}"
-        print(msg)
-        return -1, ori_audio, round(time.time() - _st, 4)
-
-def determine_pause(audio: np.ndarray, sampling_rate: int, state: VoiceAppState) -> bool:
-    """Determine if a pause happened in the audio stream"""
-    temp_audio = audio
-    dur_vad, _, time_vad = run_vad(temp_audio, sampling_rate)
-    duration = len(audio) / sampling_rate
-
-    if dur_vad > 0.5 and not state.started_talking:
-        print("started talking")
-        state.started_talking = True
-        return False
-
-    print(f"duration_after_vad: {dur_vad:.3f} s, time_vad: {time_vad:.3f} s")
-    return (duration - dur_vad) > 1
-
-def process_voice_audio(audio: tuple, state: VoiceAppState):
-    """Process streaming audio input"""
-    if not VOICE_DEPENDENCIES_AVAILABLE or voice_model is None:
-        return None, state
-        
-    if state.stream is None:
-        state.stream = audio[1]
-        state.sampling_rate = audio[0]
-    else:
-        state.stream = np.concatenate((state.stream, audio[1]))
-
-    pause_detected = determine_pause(state.stream, state.sampling_rate, state)
-    state.pause_detected = pause_detected
-
-    if state.pause_detected and state.started_talking:
-        return gr.Audio(recording=False), state
-    return None, state
-
-def generate_voice_response(state: VoiceAppState):
-    """Generate voice response from audio input"""
-    if not VOICE_DEPENDENCIES_AVAILABLE or voice_model is None:
-        return None, VoiceAppState()
-        
-    if not state.pause_detected and not state.started_talking:
-        return None, VoiceAppState()
-    
-    try:
-        audio_buffer = io.BytesIO()
-        segment = AudioSegment(
-            state.stream.tobytes(),
-            frame_rate=state.sampling_rate,
-            sample_width=state.stream.dtype.itemsize,
-            channels=(1 if len(state.stream.shape) == 1 else state.stream.shape[1]),
-        )
-        segment.export(audio_buffer, format="wav")
-
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            f.write(audio_buffer.getvalue())
-            temp_audio_path = f.name
-        
-        try:
-            # Generate text response from audio
-            text_response = voice_model.generate_response(temp_audio_path)
-            print(f"Generated voice response: {text_response}")
-            
-            # Convert text to speech
-            audio_response = voice_model.text_to_speech_simple(text_response)
-            
-            # Convert to format expected by Gradio
-            audio_segment = AudioSegment.from_file(io.BytesIO(audio_response), format="wav")
-            audio_array = np.array(audio_segment.get_array_of_samples())
-            
-            if audio_segment.channels == 2:
-                audio_array = audio_array.reshape((-1, 2))
-            
-            # Update conversation history
-            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)
-            
-        finally:
-            if os.path.exists(temp_audio_path):
-                os.unlink(temp_audio_path)
-                
-    except Exception as e:
-        print(f"Error generating voice response: {e}")
-        return None, VoiceAppState()
-
-def start_voice_recording(state: VoiceAppState):
-    """Start recording user voice input"""
-    if not state.stopped:
-        return gr.Audio(recording=True)
-    return gr.Audio(recording=False)
-
-def chunk_document(text, chunk_size=1024, overlap=100):
-    """Split document into overlapping chunks for RAG"""
-    words = text.split()
-    chunks = []
-    
-    for i in range(0, len(words), chunk_size - overlap):
-        chunk = ' '.join(words[i:i + chunk_size])
-        if chunk.strip():
-            chunks.append(chunk)
-    
-    return chunks
-
-def create_embeddings(chunks):
-    """Create embeddings for document chunks"""
-    if embedding_model is None:
-        return None
-    
-    try:
-        # Process in smaller batches on CPU
-        batch_size = 32
-        embeddings = []
-        
-        for i in range(0, len(chunks), batch_size):
-            batch = chunks[i:i + batch_size]
-            batch_embeddings = embedding_model.encode(batch, show_progress_bar=False)
-            embeddings.extend(batch_embeddings)
-        
-        return np.array(embeddings)
-    except Exception as e:
-        print(f"Error creating embeddings: {e}")
-        return None
-
-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
-    
-    try:
-        question_embedding = embedding_model.encode([question], show_progress_bar=False)
-        similarities = cosine_similarity(question_embedding, embeddings)[0]
-        
-        # Get top-k most similar chunks
-        top_indices = np.argsort(similarities)[-top_k:][::-1]
-        relevant_chunks = [chunks[i] for i in top_indices]
-        
-        return relevant_chunks
-    except Exception as e:
-        print(f"Error retrieving chunks: {e}")
-        return chunks[:3]  # Fallback
-
-def process_uploaded_pdf(pdf_file, progress=gr.Progress()):
-    """Main processing function for uploaded PDF"""
-    global processed_markdown, show_results_tab, document_chunks, document_embeddings
-    
-    if pdf_file is None:
-        return "❌ No PDF uploaded", gr.Tabs(visible=False)
-    
-    try:
-        # Process PDF
-        progress(0.1, desc="Processing PDF...")
-        combined_markdown, status = process_pdf_document(pdf_file, dolphin_model, progress)
-        
-        if status == "processing_complete":
-            processed_markdown = combined_markdown
-            
-            # Create chunks and embeddings for RAG
-            progress(0.9, desc="Creating document chunks for RAG...")
-            document_chunks = chunk_document(processed_markdown)
-            document_embeddings = create_embeddings(document_chunks)
-            print(f"Created {len(document_chunks)} chunks")
-            
-            show_results_tab = True
-            progress(1.0, desc="PDF processed successfully!")
-            return "✅ PDF processed successfully! Chatbot is ready in the Chat tab.", gr.Tabs(visible=True)
-        else:
-            show_results_tab = False
-            return combined_markdown, gr.Tabs(visible=False)
-            
-    except Exception as e:
-        show_results_tab = False
-        error_msg = f"❌ Error processing PDF: {str(e)}"
-        return error_msg, gr.Tabs(visible=False)
-
-
-def get_processed_markdown():
-    """Return the processed markdown content"""
-    global processed_markdown
-    return processed_markdown if processed_markdown else "No document processed yet."
-
-
-def clear_all():
-    """Clear all data and hide results tab"""
-    global processed_markdown, show_results_tab, document_chunks, document_embeddings
-    processed_markdown = ""
-    show_results_tab = False
-    document_chunks = []
-    document_embeddings = None
-    
-    # Clear GPU cache
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    
-    return None, "", gr.Tabs(visible=False)
-
-
-# Create Gradio interface
-with gr.Blocks(
-    title="DOLPHIN PDF AI - Local Gemma 3n", 
-    theme=gr.themes.Soft(),
-    css="""
-    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
-    
-    * {
-        font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
-    }
-    
-    .main-container { 
-        max-width: 1000px; 
-        margin: 0 auto; 
-    }
-    .upload-container { 
-        text-align: center; 
-        padding: 40px 20px;
-        border: 2px dashed #e0e0e0;
-        border-radius: 15px;
-        margin: 20px 0;
-    }
-    .upload-button {
-        font-size: 18px !important;
-        padding: 15px 30px !important;
-        margin: 20px 0 !important;
-        font-weight: 600 !important;
-    }
-    .status-message {
-        text-align: center;
-        padding: 15px;
-        margin: 10px 0;
-        border-radius: 8px;
-        font-weight: 500;
-    }
-    .chatbot-container {
-        max-height: 600px;
-    }
-    h1, h2, h3 {
-        font-weight: 700 !important;
-    }
-    #progress-container {
-        margin: 10px 0;
-        min-height: 20px;
-    }
-    """
-) as demo:
-    
-    with gr.Tabs() as main_tabs:
-        # Home Tab
-        with gr.TabItem("🏠 Home", id="home"):
-            embedding_status = "✅ RAG ready" if embedding_model else "❌ RAG not loaded"
-            voice_status = "✅ Voice chat ready" if VOICE_DEPENDENCIES_AVAILABLE and voice_model else "❌ Voice chat not available"
-            gr.Markdown(
-                "# Scholar Express - Local Gemma 3n Version with Voice\n"
-                "### Upload a research paper to get a web-friendly version with AI-generated alt text for accessibility. Includes an AI chatbot and voice chat powered by local Gemma 3n.\n"
-                f"**System:** {model_status}\n"
-                f"**RAG System:** {embedding_status}\n"
-                f"**Voice Chat:** {voice_status}\n"
-                f"**DOLPHIN:** Local model for PDF processing\n"
-                f"**Gemma 3n:** Local model for alt text generation, chat, and voice\n"
-                f"**Alt Text:** Gemma 3n generates descriptive alt text for images\n"
-                f"**GPU:** {'CUDA available' if torch.cuda.is_available() else 'CPU only'}\n\n"
-                "**Features:**\n"
-                "- 📄 PDF processing with OCR and layout analysis\n"
-                "- 💬 Text-based chat about your documents\n"
-                "- 🎙️ Voice chat with Gemma 3n (new!)\n"
-                "- ♿ AI-generated alt text for accessibility"
-            )
-            
-            with gr.Column(elem_classes="upload-container"):
-                gr.Markdown("## 📄 Upload Your PDF Document")
-                
-                pdf_input = gr.File(
-                    file_types=[".pdf"],
-                    label="",
-                    height=150,
-                    elem_id="pdf_upload"
-                )
-                
-                process_btn = gr.Button(
-                    "🚀 Process PDF", 
-                    variant="primary", 
-                    size="lg",
-                    elem_classes="upload-button"
-                )
-                
-                clear_btn = gr.Button(
-                    "🗑️ Clear", 
-                    variant="secondary"
-                )
-            
-            # Dedicated progress space
-            progress_space = gr.HTML(
-                value="",
-                visible=False,
-                elem_id="progress-container"
-            )
-            
-            # Status output (hidden during processing)
-            status_output = gr.Markdown(
-                "",
-                elem_classes="status-message"
-            )
-        
-        # Results Tab (initially hidden)
-        with gr.TabItem("📖 Document", id="results", visible=False) as results_tab:
-            gr.Markdown("## Processed Document")
-            
-            markdown_display = gr.Markdown(
-                value="",
-                latex_delimiters=[
-                    {"left": "$$", "right": "$$", "display": True},
-                    {"left": "$", "right": "$", "display": False}
-                ],
-                height=700
-            )
-        
-        # Chatbot Tab (initially hidden)
-        with gr.TabItem("💬 Chat", id="chat", visible=False) as chat_tab:
-            gr.Markdown("## Ask Questions About Your Document")
-            
-            chatbot = gr.Chatbot(
-                value=[],
-                height=500,
-                type='messages',
-                elem_classes="chatbot-container",
-                placeholder="Your conversation will appear here once you process a document..."
-            )
-            
-            with gr.Row():
-                msg_input = gr.Textbox(
-                    placeholder="Ask a question about the processed document...",
-                    scale=4,
-                    container=False
-                )
-                send_btn = gr.Button("Send", variant="primary", scale=1)
-            
-            gr.Markdown(
-                "*Ask questions about your processed document. The AI uses RAG (Retrieval-Augmented Generation) with local Gemma 3n to find relevant sections and provide accurate answers.*",
-                elem_id="chat-notice"
-            )
-        
-        # Voice Chat Tab
-        with gr.TabItem("🎙️ Talk with Gemma", id="voice") as voice_tab:
-            voice_status = "✅ Voice chat ready" if VOICE_DEPENDENCIES_AVAILABLE and voice_model else "❌ Voice chat not available"
-            gr.Markdown(f"## Voice Chat with Gemma 3n\n{voice_status}")
-            
-            if VOICE_DEPENDENCIES_AVAILABLE and voice_model:
-                with gr.Row():
-                    with gr.Column():
-                        voice_input_audio = gr.Audio(
-                            label="Speak to Gemma", 
-                            sources=["microphone"], 
-                            type="numpy",
-                            streaming=True
-                        )
-                    with gr.Column():
-                        voice_output_audio = gr.Audio(
-                            label="Gemma's Response", 
-                            streaming=True, 
-                            autoplay=True
-                        )
-                        voice_chatbot = gr.Chatbot(
-                            label="Voice Conversation", 
-                            type="messages",
-                            height=300
-                        )
-                
-                with gr.Row():
-                    voice_stop_btn = gr.Button("⏹️ Stop Conversation", variant="stop")
-                
-                gr.Markdown(
-                    "*Speak naturally to Gemma 3n. The AI will listen to your voice, process your speech, and respond with both text and voice. You can have conversations before or after processing PDFs.*"
-                )
-                
-                # Voice state
-                voice_state = gr.State(value=VoiceAppState())
-            else:
-                gr.Markdown(
-                    "### Voice chat is not available\n"
-                    "To enable voice chat, please install the required dependencies:\n"
-                    "```bash\n"
-                    "pip install librosa pydub onnxruntime\n"
-                    "```\n"
-                    "And ensure the voice_chat directory is properly set up."
-                )
-    
-    # Event handlers
-    process_btn.click(
-        fn=process_uploaded_pdf,
-        inputs=[pdf_input],
-        outputs=[status_output, results_tab],
-        show_progress=True
-    ).then(
-        fn=get_processed_markdown,
-        outputs=[markdown_display]
-    ).then(
-        fn=lambda: gr.TabItem(visible=True),
-        outputs=[chat_tab]
-    )
-    
-    clear_btn.click(
-        fn=clear_all,
-        outputs=[pdf_input, status_output, results_tab]
-    ).then(
-        fn=lambda: gr.HTML(visible=False),
-        outputs=[progress_space]
-    ).then(
-        fn=lambda: gr.TabItem(visible=False),
-        outputs=[chat_tab]
-    )
-    
-    # Chatbot functionality with local Gemma 3n
-    def chatbot_response(message, history):
-        if not message.strip():
-            return history
-        
-        if not processed_markdown:
-            return history + [{"role": "user", "content": message}, {"role": "assistant", "content": "❌ Please process a PDF document first before asking questions."}]
-        
-        try:
-            # Use RAG to get relevant chunks from markdown
-            if document_chunks and len(document_chunks) > 0:
-                relevant_chunks = retrieve_relevant_chunks(message, document_chunks, document_embeddings, top_k=3)
-                context = "\n\n".join(relevant_chunks)
-                # Smart truncation: aim for ~6000 chars for local model
-                if len(context) > 6000:
-                    # Try to cut at sentence boundaries
-                    sentences = context[:6000].split('.')
-                    context = '.'.join(sentences[:-1]) + '...' if len(sentences) > 1 else context[:6000] + '...'
-            else:
-                # Fallback to truncated document if RAG fails
-                context = processed_markdown[:6000] + "..." if len(processed_markdown) > 6000 else processed_markdown
-            
-            # Create prompt for Gemma 3n
-            prompt = f"""You are a helpful assistant that answers questions about documents. Use the provided context to answer questions accurately and concisely.
-
-Context from the document:
-{context}
-
-Question: {message}
-
-Please provide a clear and helpful answer based on the context provided."""
-            
-            # Generate response using local Gemma 3n
-            response_text = gemma_model.chat(prompt)
-            return history + [{"role": "user", "content": message}, {"role": "assistant", "content": response_text}]
-            
-        except Exception as e:
-            error_msg = f"❌ Error generating response: {str(e)}"
-            print(f"Full error: {e}")
-            import traceback
-            traceback.print_exc()
-            return history + [{"role": "user", "content": message}, {"role": "assistant", "content": error_msg}]
-    
-    send_btn.click(
-        fn=chatbot_response,
-        inputs=[msg_input, chatbot],
-        outputs=[chatbot]
-    ).then(
-        lambda: "",
-        outputs=[msg_input]
-    )
-    
-    # Also allow Enter key to send message
-    msg_input.submit(
-        fn=chatbot_response,
-        inputs=[msg_input, chatbot],
-        outputs=[chatbot]
-    ).then(
-        lambda: "",
-        outputs=[msg_input]
-    )
-    
-    # Voice chat event handlers
-    if VOICE_DEPENDENCIES_AVAILABLE and voice_model:
-        # Stream processing
-        voice_stream = voice_input_audio.stream(
-            process_voice_audio,
-            [voice_input_audio, voice_state],
-            [voice_input_audio, voice_state],
-            stream_every=0.50,
-            time_limit=30,
-        )
-        
-        # Response generation
-        voice_respond = voice_input_audio.stop_recording(
-            generate_voice_response,
-            [voice_state],
-            [voice_output_audio, voice_state]
-        )
-        
-        # Update chatbot display
-        voice_respond.then(
-            lambda s: s.conversation, 
-            [voice_state], 
-            [voice_chatbot]
-        )
-        
-        # Restart recording
-        voice_restart = voice_output_audio.stop(
-            start_voice_recording,
-            [voice_state],
-            [voice_input_audio]
-        )
-        
-        # Stop conversation
-        voice_stop_btn.click(
-            lambda: (VoiceAppState(stopped=True), gr.Audio(recording=False)), 
-            None,
-            [voice_state, voice_input_audio], 
-            cancels=[voice_respond, voice_restart]
-        )
-
-
-if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True,
-        max_threads=4,
-        inbrowser=False,
-        quiet=True
-    )

chat.py

@@ -1,198 +0,0 @@
-""" 
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
-import os
-import warnings
-from collections import OrderedDict
-
-from omegaconf import ListConfig
-
-warnings.filterwarnings("ignore", category=UserWarning)
-warnings.filterwarnings("ignore", category=FutureWarning)
-os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python")
-
-import torch
-from PIL import Image
-from transformers import PreTrainedTokenizerFast
-
-from utils.model import DonutConfig, DonutModel, SwinEncoder
-from utils.processor import DolphinProcessor
-
-
-def try_rename_lagacy_weights(ckpt, output_path=""):
-    if "state_dict" in ckpt.keys():
-        ckpt = ckpt["state_dict"]
-    if "module" in ckpt.keys():
-        ckpt = ckpt["module"]
-    new_ckpt = OrderedDict()
-    for k, v in ckpt.items():
-        if k.startswith("model."):
-            k = k[len("model.") :]
-        if k.startswith("encoder"):
-            new_ckpt["vpm" + k[len("encoder") :]] = v
-        elif k.startswith("decoder"):
-            new_ckpt["llm" + k[len("encoder") :]] = v
-        else:
-            new_ckpt[k] = v
-    if output_path:
-        torch.save(new_ckpt, output_path)
-    return new_ckpt
-
-
-def convert_listconfig_to_list(config):
-    new_config = {}
-    for k, v in config.items():
-        if isinstance(v, ListConfig):
-            new_config[k] = list(v)
-        else:
-            new_config[k] = v
-    return new_config
-
-
-class DOLPHIN:
-    def __init__(self, config, ckpt_path="") -> None:
-        self.model_args = config.model
-        self.swin_args = config.model.pop("swin_args")
-        self.swin_args = convert_listconfig_to_list(self.swin_args)
-
-        vision_tower = SwinEncoder(
-            input_size=self.swin_args["img_size"],
-            patch_size=self.swin_args["patch_size"],
-            embed_dim=self.swin_args["embed_dim"],
-            window_size=self.swin_args["window_size"],
-            encoder_layer=self.swin_args["encoder_layer"],
-            num_heads=self.swin_args["num_heads"],
-            align_long_axis=self.swin_args["align_long_axis"],
-        )
-
-        self.tokenizer = PreTrainedTokenizerFast(tokenizer_file=self.model_args.tokenizer_path)
-        self.tokenizer.pad_token = "<pad>"
-        self.tokenizer.bos_token = "<s>"
-        self.tokenizer.eos_token = "</s>"
-        self.tokenizer.unk_token = "<unk>"
-
-        if self.model_args.get("extra_answer_tokens", False):
-            # print("Allowing multitask training: adding <Answer/> to the tokenizer.")
-            prompt_end_token = " <Answer/>"
-            self.tokenizer.add_special_tokens({"additional_special_tokens": sorted(set([prompt_end_token]))})
-            self.tokenizer._prompt_end_token = prompt_end_token
-            self.tokenizer._prompt_end_token_id = self.tokenizer.convert_tokens_to_ids(prompt_end_token)
-
-        donut_config = DonutConfig(
-            decoder_layer=self.model_args.decoder_layer,
-            max_length=self.model_args.max_length,
-            max_position_embeddings=self.model_args.max_position_embeddings,
-            hidden_dimension=self.model_args.hidden_dimension,
-        )
-
-        self.model = DonutModel(config=donut_config, vision_tower=vision_tower, tokenizer=self.tokenizer)
-        if self.model_args.model_name_or_path:
-            ckpt = torch.load(self.model_args.model_name_or_path)
-            ckpt = try_rename_lagacy_weights(ckpt)
-            self.model.load_state_dict(ckpt, strict=True)
-
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(device)
-        self.model.eval()
-        transform_args = {
-            "input_size": self.swin_args["img_size"],
-            "max_length": self.model_args.max_length,
-        }
-        self.processor = DolphinProcessor({}, self.tokenizer, transform_args=transform_args)
-
-    def chat(
-        self,
-        question,
-        image,
-        return_raw=False,
-        return_score=False,
-        return_img_size=False,
-        only_return_img_size=False,
-        max_batch_size=16,
-    ):
-
-        def _preprocess_image(image):
-            if isinstance(image, str):
-                image = Image.open(image).convert("RGB")
-            if return_img_size or only_return_img_size:
-                image_tensor, ori_size = self.processor.process_image_for_inference(image, return_img_size=True)
-            else:
-                image_tensor = self.processor.process_image_for_inference(image, return_img_size=False)
-                ori_size = None
-            return image_tensor, ori_size
-
-        def _preprocess_prompt(question):
-            if self.model_args.get("extra_answer_tokens", False):
-                if self.tokenizer._prompt_end_token not in question:
-                    question = question + self.tokenizer._prompt_end_token
-            prompt_ids = self.processor.process_prompt_for_inference(question)
-            return prompt_ids
-
-        def _preprocess_prompt_batch(question):
-            if self.model_args.get("extra_answer_tokens", False):
-                for i in range(len(question)):
-                    if self.tokenizer._prompt_end_token not in question[i]:
-                        question[i] = question[i] + self.tokenizer._prompt_end_token
-                    if not question[i].startswith("<s>"):
-                        question[i] = "<s>" + question[i]
-            return question
-
-        def _postprocess(output, question):
-            output = output.replace("<s>", "").replace(question, "").replace("</s>", "").replace("<pad>", "")
-            if self.model_args.get("extra_answer_tokens", False):
-                output = output.split(self.tokenizer._prompt_end_token)[-1]
-            return output
-
-        if isinstance(question, list):
-            image_tensor_list = []
-            for i in image:
-                image_tensor, ori_size = _preprocess_image(i)
-                image_tensor_list.append(image_tensor)
-            image_tensor = torch.cat(image_tensor_list, dim=0)
-
-            question = _preprocess_prompt_batch(question)
-            self.processor.tokenizer.padding_side = "left"
-            prompt_ids = self.processor.tokenizer(
-                question, add_special_tokens=False, return_tensors="pt", padding=True
-            ).input_ids
-        else:
-            image_tensor, ori_size = _preprocess_image(image)
-            prompt_ids = _preprocess_prompt(question)
-
-        if only_return_img_size:
-            return ori_size
-
-        model_output_batch = []
-        for i in range(0, image_tensor.shape[0], max_batch_size):
-            image_tensor_batch = image_tensor[i : i + max_batch_size]
-            prompt_ids_batch = prompt_ids[i : i + max_batch_size]
-            model_output = self.model.inference(image_tensors=image_tensor_batch, prompt_ids=prompt_ids_batch)
-            model_output_batch.append(model_output)
-        model_output = {}
-        for k, v in model_output_batch[0].items():
-            if isinstance(v, torch.Tensor):
-                model_output[k] = sum(
-                    [v_batch[k].cpu().numpy().tolist() for v_batch in model_output_batch],
-                    [],
-                )
-            else:
-                model_output[k] = sum([v_batch[k] for v_batch in model_output_batch], [])
-
-        if return_raw:
-            if return_img_size:
-                return model_output, ori_size
-            return model_output
-        else:
-            if isinstance(question, list):
-                output = [_postprocess(model_output["repetitions"][i], question[i]) for i in range(len(question))]
-                score = model_output["scores"]
-            else:
-                output = _postprocess(model_output["repetitions"][0], question)
-                score = model_output["scores"][0]
-            if return_score:
-                return output, score
-            if return_img_size:
-                return output, ori_size
-            return output

config/Dolphin.yaml

@@ -1,17 +0,0 @@
-model:
-  model_name_or_path: "./checkpoints/dolphin_model.bin"
-  tokenizer_path: "./checkpoints/dolphin_tokenizer.json"
-  extra_answer_tokens: True   # add <Answer/> token
-  max_length: 4096
-  decoder_layer: 10
-  max_position_embeddings: 4096
-  hidden_dimension: 1024
-  swin_args:
-    name: 'swin'
-    img_size: [896, 896]
-    patch_size: 4
-    embed_dim: 128
-    align_long_axis: False
-    window_size: 7
-    encoder_layer: [2, 2, 14, 2]
-    num_heads: [4, 8, 16, 32]

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,195 +0,0 @@
-"""
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
-import argparse
-import glob
-import os
-
-import torch
-from PIL import Image
-from transformers import AutoProcessor, VisionEncoderDecoderModel
-
-from utils.utils import *
-
-
-class DOLPHIN:
-    def __init__(self, model_id_or_path):
-        """Initialize the Hugging Face model
-        
-        Args:
-            model_id_or_path: Path to local model or Hugging Face model ID
-        """
-        # Load model from local path or Hugging Face hub
-        self.processor = AutoProcessor.from_pretrained(model_id_or_path)
-        self.model = VisionEncoderDecoderModel.from_pretrained(model_id_or_path)
-        self.model.eval()
-        
-        # Set device and precision
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(self.device)
-        self.model = self.model.half()  # Always use half precision by default
-        
-        # set tokenizer
-        self.tokenizer = self.processor.tokenizer
-        
-    def chat(self, prompt, image):
-        """Process an image with the given prompt
-        
-        Args:
-            prompt: Text prompt to guide the model
-            image: PIL Image to process
-            
-        Returns:
-            Generated text from the model
-        """
-        # Prepare image
-        pixel_values = self.processor(image, return_tensors="pt").pixel_values
-        pixel_values = pixel_values.half()
-            
-        # Prepare prompt
-        prompt = f"<s>{prompt} <Answer/>"
-        prompt_ids = self.tokenizer(
-            prompt, 
-            add_special_tokens=False, 
-            return_tensors="pt"
-        ).input_ids.to(self.device)
-        
-        decoder_attention_mask = torch.ones_like(prompt_ids)
-        
-        # Generate text
-        outputs = self.model.generate(
-            pixel_values=pixel_values.to(self.device),
-            decoder_input_ids=prompt_ids,
-            decoder_attention_mask=decoder_attention_mask,
-            min_length=1,
-            max_length=4096,
-            pad_token_id=self.tokenizer.pad_token_id,
-            eos_token_id=self.tokenizer.eos_token_id,
-            use_cache=True,
-            bad_words_ids=[[self.tokenizer.unk_token_id]],
-            return_dict_in_generate=True,
-            do_sample=False,
-            num_beams=1,
-            repetition_penalty=1.1,
-            temperature=1.0
-        )
-        
-        # Process the output
-        sequence = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)[0]
-        sequence = sequence.replace(prompt, "").replace("<pad>", "").replace("</s>", "").strip()
-        
-        return sequence
-
-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: HFModel 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("--model_path", default="./hf_model", help="Path to Hugging Face model")
-    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
-    model = DOLPHIN(args.model_path)
-    
-    # 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_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,365 +0,0 @@
-""" 
-Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-SPDX-License-Identifier: MIT
-"""
-
-import argparse
-import glob
-import os
-
-import cv2
-import torch
-from PIL import Image
-from transformers import AutoProcessor, VisionEncoderDecoderModel
-
-from utils.utils import *
-
-
-class DOLPHIN:
-    def __init__(self, model_id_or_path):
-        """Initialize the Hugging Face model
-        
-        Args:
-            model_id_or_path: Path to local model or Hugging Face model ID
-        """
-        # Load model from local path or Hugging Face hub
-        self.processor = AutoProcessor.from_pretrained(model_id_or_path)
-        self.model = VisionEncoderDecoderModel.from_pretrained(model_id_or_path)
-        self.model.eval()
-        
-        # Set device and precision
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(self.device)
-        self.model = self.model.half()  # Always use half precision by default
-        
-        # set tokenizer
-        self.tokenizer = self.processor.tokenizer
-        
-    def chat(self, prompt, image):
-        """Process an image or batch of images with the given prompt(s)
-        
-        Args:
-            prompt: Text prompt or list of prompts to guide the model
-            image: PIL Image or list of PIL Images to process
-            
-        Returns:
-            Generated text or list of texts from the model
-        """
-        # Check if we're dealing with a batch
-        is_batch = isinstance(image, list)
-        
-        if not is_batch:
-            # Single image, wrap it in a list for consistent processing
-            images = [image]
-            prompts = [prompt]
-        else:
-            # Batch of images
-            images = image
-            prompts = prompt if isinstance(prompt, list) else [prompt] * len(images)
-        
-        # Prepare image
-        batch_inputs = self.processor(images, return_tensors="pt", padding=True)
-        batch_pixel_values = batch_inputs.pixel_values.half().to(self.device)
-        
-        # Prepare prompt
-        prompts = [f"<s>{p} <Answer/>" for p in prompts]
-        batch_prompt_inputs = self.tokenizer(
-            prompts,
-            add_special_tokens=False,
-            return_tensors="pt"
-        )
-
-        batch_prompt_ids = batch_prompt_inputs.input_ids.to(self.device)
-        batch_attention_mask = batch_prompt_inputs.attention_mask.to(self.device)
-        
-        # Generate text
-        outputs = self.model.generate(
-            pixel_values=batch_pixel_values,
-            decoder_input_ids=batch_prompt_ids,
-            decoder_attention_mask=batch_attention_mask,
-            min_length=1,
-            max_length=4096,
-            pad_token_id=self.tokenizer.pad_token_id,
-            eos_token_id=self.tokenizer.eos_token_id,
-            use_cache=True,
-            bad_words_ids=[[self.tokenizer.unk_token_id]],
-            return_dict_in_generate=True,
-            do_sample=False,
-            num_beams=1,
-            repetition_penalty=1.1,
-            temperature=1.0
-        )
-        
-        # Process output
-        sequences = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
-        
-        # Clean prompt text from output
-        results = []
-        for i, sequence in enumerate(sequences):
-            cleaned = sequence.replace(prompts[i], "").replace("<pad>", "").replace("</s>", "").strip()
-            results.append(cleaned)
-            
-        # Return a single result for single image input
-        if not is_batch:
-            return results[0]
-        return results
-
-
-def process_document(document_path, model, save_dir, max_batch_size=None):
-    """Parse documents with two stages - 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=None, 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)
-
-    # Store text and table elements separately
-    text_elements = []  # Text elements
-    table_elements = []  # Table elements
-    figure_results = []  # Image elements (no processing needed)
-    previous_box = None
-    reading_order = 0
-
-    # Collect elements to process and group by type
-    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:
-                    # Prepare element for parsing
-                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
-                    element_info = {
-                        "crop": pil_crop,
-                        "label": label,
-                        "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
-                        "reading_order": reading_order,
-                    }
-                    
-                    # Group by type
-                    if label == "tab":
-                        table_elements.append(element_info)
-                    else:  # Text elements
-                        text_elements.append(element_info)
-
-            reading_order += 1
-
-        except Exception as e:
-            print(f"Error processing bbox with label {label}: {str(e)}")
-            continue
-
-    # Initialize results list
-    recognition_results = figure_results.copy()
-    
-    # Process text elements (in batches)
-    if text_elements:
-        text_results = process_element_batch(text_elements, model, "Read text in the image.", max_batch_size)
-        recognition_results.extend(text_results)
-    
-    # Process table elements (in batches)
-    if table_elements:
-        table_results = process_element_batch(table_elements, model, "Parse the table in the image.", max_batch_size)
-        recognition_results.extend(table_results)
-
-    # Sort elements by reading order
-    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
-
-    return recognition_results
-
-
-def process_element_batch(elements, model, prompt, max_batch_size=None):
-    """Process elements of the same type in batches"""
-    results = []
-    
-    # Determine batch size
-    batch_size = len(elements)
-    if max_batch_size is not None and max_batch_size > 0:
-        batch_size = min(batch_size, max_batch_size)
-    
-    # Process in batches
-    for i in range(0, len(elements), batch_size):
-        batch_elements = elements[i:i+batch_size]
-        crops_list = [elem["crop"] for elem in batch_elements]
-        
-        # Use the same prompt for all elements in the batch
-        prompts_list = [prompt] * len(crops_list)
-        
-        # Batch inference
-        batch_results = model.chat(prompts_list, crops_list)
-        
-        # Add results
-        for j, result in enumerate(batch_results):
-            elem = batch_elements[j]
-            results.append({
-                "label": elem["label"],
-                "bbox": elem["bbox"],
-                "text": result.strip(),
-                "reading_order": elem["reading_order"],
-            })
-    
-    return results
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Document parsing based on DOLPHIN")
-    parser.add_argument("--model_path", default="./hf_model", help="Path to Hugging Face model")
-    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=16,
-        help="Maximum number of document elements to parse in a single batch (default: 16)",
-    )
-    args = parser.parse_args()
-
-    # Load Model
-    model = DOLPHIN(args.model_path)
-
-    # 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()

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()