Update app.py

app.py

@@ -1,364 +1,302 @@
-import gradio as gr
+import subprocess
+subprocess.run(
+    'pip install flash-attn==2.7.0.post2 --no-build-isolation',
+    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+    shell=True
+)
+subprocess.run(
+    'pip install transformers',
+    shell=True
+)
+
+
 import spaces
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
-from qwen_vl_utils import process_vision_info
-import torch
-from PIL import Image
 import os
-import uuid
-import io
+import re
+import logging
+from typing import List
 from threading import Thread
-from reportlab.lib.pagesizes import A4
-from reportlab.lib.styles import getSampleStyleSheet
-from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer
-from reportlab.lib.units import inch
-from reportlab.pdfbase import pdfmetrics
-from reportlab.pdfbase.ttfonts import TTFont
-import docx
-from docx.enum.text import WD_ALIGN_PARAGRAPH
-
-# Define model options
-MODEL_OPTIONS = {
-    "Qwen2VL Base": "Qwen/Qwen2-VL-2B-Instruct",
-    "Latex OCR": "prithivMLmods/Qwen2-VL-OCR-2B-Instruct",
-    "Math Prase": "prithivMLmods/Qwen2-VL-Math-Prase-2B-Instruct",
-    "Text Analogy Ocrtest": "prithivMLmods/Qwen2-VL-Ocrtest-2B-Instruct"
-}
+import base64
+
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+# ----------------------------------------------------------------------
+# 1. Setup Model & Tokenizer
+# ----------------------------------------------------------------------
+model_name = 'prithivMLmods/Raptor-X5-UIGEN'  # Change as needed
+use_thread = True  # Generation happens in a background thread
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True
+).to("cuda")
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+logging.getLogger("httpx").setLevel(logging.WARNING)
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ----------------------------------------------------------------------
+# 2. Two-Phase Prompt Templates
+# ----------------------------------------------------------------------
+s1_inference_prompt_think_only = """<|im_start|>user
+{question}<|im_end|>
+<|im_start|>assistant
+<|im_start|>think
+"""
 
-# Preload models and processors into CUDA
-models = {}
-processors = {}
-for name, model_id in MODEL_OPTIONS.items():
-    print(f"Loading {name}...")
-    models[name] = Qwen2VLForConditionalGeneration.from_pretrained(
-        model_id,
-        trust_remote_code=True,
-        torch_dtype=torch.float16
-    ).to("cuda").eval()
-    processors[name] = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-
-# Get valid image extensions from PIL
-image_extensions = Image.registered_extensions()
-
-def identify_and_save_blob(blob_path):
-    """Identifies if the blob is an image and saves it."""
-    try:
-        with open(blob_path, 'rb') as file:
-            blob_content = file.read()
-            try:
-                Image.open(io.BytesIO(blob_content)).verify()  # Check if it's a valid image
-                extension = ".png"  # Default to PNG for saving
-                media_type = "image"
-            except (IOError, SyntaxError):
-                raise ValueError("Unsupported media type. Please upload a valid image.")
-
-            filename = f"temp_{uuid.uuid4()}_media{extension}"
-            with open(filename, "wb") as f:
-                f.write(blob_content)
-
-            return filename, media_type
-
-    except FileNotFoundError:
-        raise ValueError(f"The file {blob_path} was not found.")
-    except Exception as e:
-        raise ValueError(f"An error occurred while processing the file: {e}")
-
-def get_media_file(media_input):
+# ----------------------------------------------------------------------
+# 3. Generation Parameter Setup
+# ----------------------------------------------------------------------
+THINK_MAX_NEW_TOKENS = 12000
+ANSWER_MAX_NEW_TOKENS = 12000
+
+def initialize_gen_kwargs():
+    return {
+        "max_new_tokens": 1024,  # default; will be overwritten per phase
+        "do_sample": True,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "repetition_penalty": 1.05,
+        # "eos_token_id": model.generation_config.eos_token_id,  # Removed to avoid premature stopping
+        "pad_token_id": tokenizer.pad_token_id,
+        "use_cache": True,
+        "streamer": None  # dynamically added
+    }
+
+# ----------------------------------------------------------------------
+# 4. Helper to submit chat
+# ----------------------------------------------------------------------
+def submit_chat(chatbot, text_input):
+    if not text_input.strip():
+        return chatbot, ""
+    response = ""
+    chatbot.append((text_input, response))
+    return chatbot, ""
+
+# ----------------------------------------------------------------------
+# 5. Artifacts Handling
+#    We parse code from the final answer and display it in an iframe
+# ----------------------------------------------------------------------
+def extract_html_code_block(text: str) -> str:
     """
-    Ensures that the media input is a file path.
-    If it is a PIL image, it saves it temporarily and returns the file path.
+    Look for a ```html ... ``` block in the text. 
+    If found, return only that block content. 
+    Otherwise, return the entire text.
     """
-    if isinstance(media_input, str):
-        return media_input  # Already a file path
+    pattern = r'```html\s*(.*?)\s*```'
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
     else:
-        if not isinstance(media_input, Image.Image):
-            # Convert numpy array to PIL image if needed
-            media_input = Image.fromarray(media_input)
-        temp_filename = f"temp_{uuid.uuid4()}.png"
-        media_input.save(temp_filename)
-        return temp_filename
+        return text.strip()
 
+def send_to_sandbox(html_code: str) -> str:
+    """
+    Convert the code to a data URI iframe so it can be rendered
+    inside Gradio HTML component.
+    """
+    encoded_html = base64.b64encode(html_code.encode('utf-8')).decode('utf-8')
+    data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
+    return f'<iframe src="{data_uri}" width="100%" height="920px"></iframe>'
+
+# ----------------------------------------------------------------------
+# 6. The Two-Phase Streaming Inference
+#    - Phase 1: "think" (chain-of-thought)
+#    - Phase 2: "answer"
+# ----------------------------------------------------------------------
 @spaces.GPU
-def qwen_inference(model_name, media_input, text_input=None):
-    """Handles inference for the selected model."""
-    model = models[model_name]
-    processor = processors[model_name]
-
-    # Determine media type and obtain a file path if needed
-    if isinstance(media_input, str):
-        media_path = media_input
-        if media_path.endswith(tuple(image_extensions.keys())):
-            media_type = "image"
-        else:
-            try:
-                media_path, media_type = identify_and_save_blob(media_input)
-            except Exception as e:
-                raise ValueError("Unsupported media type. Please upload a valid image.")
-    else:
-        # media_input is a PIL image (or numpy array) coming from gr.Image
-        media_path = get_media_file(media_input)
-        media_type = "image"
-
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": media_type,
-                    media_type: media_path
-                },
-                {"type": "text", "text": text_input},
-            ],
-        }
-    ]
-
-    text = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    image_inputs, _ = process_vision_info(messages)
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        padding=True,
-        return_tensors="pt",
-    ).to("cuda")
-
-    streamer = TextIteratorStreamer(
-        processor.tokenizer, skip_prompt=True, skip_special_tokens=True
-    )
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        # Remove <|im_end|> or similar tokens from the output
-        buffer = buffer.replace("<|im_end|>", "")
-        yield buffer
-
-def format_plain_text(output_text):
-    """Formats the output text as plain text without LaTeX delimiters."""
-    plain_text = output_text.replace("\\(", "").replace("\\)", "").replace("\\[", "").replace("\\]", "")
-    return plain_text
-
-def generate_document(media_input, output_text, file_format, font_choice, font_size, line_spacing, alignment, image_size):
-    """Generates a document with the input image and plain text output."""
-    # Ensure media_input is a file path.
-    media_path = get_media_file(media_input)
-    plain_text = format_plain_text(output_text)
-    if file_format == "pdf":
-        return generate_pdf(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size)
-    elif file_format == "docx":
-        return generate_docx(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size)
-
-def generate_pdf(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size):
-    """Generates a PDF document."""
-    filename = f"output_{uuid.uuid4()}.pdf"
-    doc = SimpleDocTemplate(
-        filename,
-        pagesize=A4,
-        rightMargin=inch,
-        leftMargin=inch,
-        topMargin=inch,
-        bottomMargin=inch
-    )
-    styles = getSampleStyleSheet()
-    styles["Normal"].fontName = font_choice
-    styles["Normal"].fontSize = int(font_size)
-    styles["Normal"].leading = int(font_size) * line_spacing
-    styles["Normal"].alignment = {
-        "Left": 0,
-        "Center": 1,
-        "Right": 2,
-        "Justified": 4
-    }[alignment]
-
-    # Register font (assumes font files are available in a folder named "font")
-    font_path = f"font/{font_choice}"
-    pdfmetrics.registerFont(TTFont(font_choice, font_path))
-
-    story = []
-
-    # Add image with size adjustment
-    image_sizes = {
-        "Small": (200, 200),
-        "Medium": (400, 400),
-        "Large": (600, 600)
+def ovis_chat(chatbot: List[List[str]]):
+    # Phase 1: chain-of-thought
+    last_query = chatbot[-1][0]
+    formatted_think_prompt = s1_inference_prompt_think_only.format(question=last_query)
+    input_ids_think = tokenizer.encode(formatted_think_prompt, return_tensors="pt").to(model.device)
+    attention_mask_think = torch.ne(input_ids_think, tokenizer.pad_token_id).to(model.device)
+
+    think_inputs = {
+        "input_ids": input_ids_think,
+        "attention_mask": attention_mask_think
     }
-    img = RLImage(media_path, width=image_sizes[image_size][0], height=image_sizes[image_size][1])
-    story.append(img)
-    story.append(Spacer(1, 12))
-
-    # Add plain text output
-    text_para = Paragraph(plain_text, styles["Normal"])
-    story.append(text_para)
-
-    doc.build(story)
-    return filename
-
-def generate_docx(media_path, plain_text, font_choice, font_size, line_spacing, alignment, image_size):
-    """Generates a DOCX document."""
-    filename = f"output_{uuid.uuid4()}.docx"
-    doc = docx.Document()
-
-    # Add image with size adjustment
-    image_sizes = {
-        "Small": docx.shared.Inches(2),
-        "Medium": docx.shared.Inches(4),
-        "Large": docx.shared.Inches(6)
+    gen_kwargs_think = initialize_gen_kwargs()
+    gen_kwargs_think["max_new_tokens"] = THINK_MAX_NEW_TOKENS
+    think_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs_think["streamer"] = think_streamer
+
+    full_think = ""
+    with torch.inference_mode():
+        thread_think = Thread(target=lambda: model.generate(**think_inputs, **gen_kwargs_think))
+        thread_think.start()
+        for new_text in think_streamer:
+            full_think += new_text
+            display_text = f"<|im_start|>think\n{full_think.strip()}"
+            chatbot[-1][1] = display_text
+            yield chatbot, ""  # second return is artifact placeholder
+        thread_think.join()
+
+    # Phase 2: answer
+    new_prompt = formatted_think_prompt + full_think.strip() + "\n<|im_start|>answer\n"
+    input_ids_answer = tokenizer.encode(new_prompt, return_tensors="pt").to(model.device)
+    attention_mask_answer = torch.ne(input_ids_answer, tokenizer.pad_token_id).to(model.device)
+
+    answer_inputs = {
+        "input_ids": input_ids_answer,
+        "attention_mask": attention_mask_answer
     }
-    doc.add_picture(media_path, width=image_sizes[image_size])
-    doc.add_paragraph()
-
-    # Add plain text output
-    paragraph = doc.add_paragraph()
-    paragraph.paragraph_format.line_spacing = line_spacing
-    paragraph.paragraph_format.alignment = {
-        "Left": WD_ALIGN_PARAGRAPH.LEFT,
-        "Center": WD_ALIGN_PARAGRAPH.CENTER,
-        "Right": WD_ALIGN_PARAGRAPH.RIGHT,
-        "Justified": WD_ALIGN_PARAGRAPH.JUSTIFY
-    }[alignment]
-    run = paragraph.add_run(plain_text)
-    run.font.name = font_choice
-    run.font.size = docx.shared.Pt(int(font_size))
-
-    doc.save(filename)
-    return filename
-
-# CSS for output styling
-css = """
-  #output {
-    height: 400px; 
-    overflow: auto; 
-    border: 1px solid #ccc; 
-  }
-.submit-btn {
-    background-color: #cf3434  !important;
-    color: white !important;
+    gen_kwargs_answer = initialize_gen_kwargs()
+    gen_kwargs_answer["max_new_tokens"] = ANSWER_MAX_NEW_TOKENS
+    answer_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    gen_kwargs_answer["streamer"] = answer_streamer
+
+    full_answer = ""
+    with torch.inference_mode():
+        thread_answer = Thread(target=lambda: model.generate(**answer_inputs, **gen_kwargs_answer))
+        thread_answer.start()
+        for new_text in answer_streamer:
+            full_answer += new_text
+            display_text = (
+                f"<|im_start|>think\n{full_think.strip()}\n\n"
+                f"<|im_start|>answer\n{full_answer.strip()}"
+            )
+            chatbot[-1][1] = display_text
+            yield chatbot, ""
+        thread_answer.join()
+
+    log_conversation(chatbot)
+
+    # Once final answer is complete, parse out HTML code block and
+    # return it as an artifact (iframe).
+    html_code = extract_html_code_block(full_answer)
+    sandbox_iframe = send_to_sandbox(html_code)
+    yield chatbot, sandbox_iframe
+
+# ----------------------------------------------------------------------
+# 7. Logging and Clearing
+# ----------------------------------------------------------------------
+def log_conversation(chatbot: List[List[str]]):
+    logger.info("[CONVERSATION]")
+    for i, (query, response) in enumerate(chatbot, 1):
+        logger.info(f"Q{i}: {query}\nA{i}: {response}")
+
+def clear_chat():
+    return [], "", ""
+
+# ----------------------------------------------------------------------
+# 8. Gradio UI Setup
+# ----------------------------------------------------------------------
+css_code = """
+.left_header {
+  display: flex;
+  flex-direction: column;
+  justify-content: center;
+  align-items: center;
+}
+
+.right_panel {
+  margin-top: 16px;
+  border: 1px solid #BFBFC4;
+  border-radius: 8px;
+  overflow: hidden;
+}
+
+.render_header {
+  height: 30px;
+  width: 100%;
+  padding: 5px 16px;
+  background-color: #f5f5f5;
+}
+
+.header_btn {
+  display: inline-block;
+  height: 10px;
+  width: 10px;
+  border-radius: 50%;
+  margin-right: 4px;
+}
+
+.render_header > .header_btn:nth-child(1) {
+  background-color: #f5222d;
 }
-.submit-btn:hover {
-    background-color: #ff2323 !important;
+
+.render_header > .header_btn:nth-child(2) {
+  background-color: #faad14;
+}
+.render_header > .header_btn:nth-child(3) {
+  background-color: #52c41a;
 }
-.download-btn {
-    background-color: #35a6d6 !important;
-    color: white !important;
+
+.right_content {
+  height: 920px; 
+  display: flex; 
+  flex-direction: column;
+  justify-content: center;
+  align-items: center;
 }
-.download-btn:hover {
-    background-color: #22bcff !important;
+
+.html_content {
+  width: 100%;
+  height: 920px;
 }
 """
 
-# Gradio app setup
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Qwen2VL: Compact Vision & Language Processing")
-
-    with gr.Tab(label="Image Input"):
-        with gr.Row():
-            with gr.Column():
-                model_choice = gr.Dropdown(
-                    label="Model Selection",
-                    choices=list(MODEL_OPTIONS.keys()),
-                    value="Latex OCR"
-                )
-                # Using gr.Image instead of gr.File for image upload
-                input_media = gr.Image(
-                    label="Upload Image", type="pil"
-                )
-                text_input = gr.Textbox(label="Question", placeholder="Ask a question about the image...")
-                submit_btn = gr.Button(value="Submit", elem_classes="submit-btn")
-            with gr.Column():
-                output_text = gr.Textbox(label="Output Text", lines=10)
-                plain_text_output = gr.Textbox(label="Standardized Plain Text", lines=10)
-
-        submit_btn.click(
-            qwen_inference, [model_choice, input_media, text_input], [output_text]
-        ).then(
-            lambda output_text: format_plain_text(output_text), [output_text], [plain_text_output]
-        )
-
-        # Add examples directly usable by clicking
-        with gr.Row():
-            gr.Examples(
-                examples=[
-                    ["examples/4.png", "solve the problem", "Math Prase"],
-                    ["examples/1.png", "summarize the letter", "Text Analogy Ocrtest"],
-                    ["examples/2.jpg", "Summarize the full image in detail", "Latex OCR"],
-                    ["examples/3.png", "Describe the photo", "Qwen2VL Base"],
-                ],
-                inputs=[input_media, text_input, model_choice],
-                outputs=[output_text, plain_text_output],
-                fn=lambda img, question, model: qwen_inference(model, img, question),
-                cache_examples=False,
+svg_content = """
+<svg width="40" height="40" viewBox="0 0 45 45" fill="none" xmlns="http://www.w3.org/2000/svg">
+  <circle cx="22.5" cy="22.5" r="22.5" fill="#5572F9"/>
+  <path d="M22.5 11.25L26.25 16.875H18.75L22.5 11.25Z" fill="white"/>
+  <path d="M22.5 33.75L26.25 28.125H18.75L22.5 33.75Z" fill="white"/>
+  <path d="M28.125 22.5L22.5 28.125L16.875 22.5L22.5 16.875L28.125 22.5Z" fill="white"/>
+</svg>
+"""
+
+with gr.Blocks(title=model_name.split('/')[-1], css=css_code) as demo:
+    gr.HTML(f"""
+        <div class="left_header" style="margin-bottom: 20px;">
+            {svg_content}
+            <h1>{model_name.split('/')[-1]} - Chat + Artifacts</h1>
+            <p>(Two-phase chain-of-thought with artifact extraction)</p>
+        </div>
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(
+                label="Chat",
+                height=520,
+                show_copy_button=True
             )
-        with gr.Row():
-            with gr.Column():
-                line_spacing = gr.Dropdown(
-                    choices=[0.5, 1.0, 1.15, 1.5, 2.0, 2.5, 3.0],
-                    value=1.5,
-                    label="Line Spacing"
-                )
-                font_size = gr.Dropdown(
-                    choices=["8", "10", "12", "14", "16", "18", "20", "22", "24"],
-                    value="12",
-                    label="Font Size"
+            with gr.Row():
+                text_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter your query...",
+                    lines=1
                 )
-                font_choice = gr.Dropdown(
-                    choices=[
-                        "DejaVuMathTeXGyre.ttf", 
-                        "FiraCode-Medium.ttf", 
-                        "InputMono-Light.ttf",
-                        "JetBrainsMono-Thin.ttf", 
-                        "ProggyCrossed Regular Mac.ttf", 
-                        "SourceCodePro-Black.ttf", 
-                        "arial.ttf", 
-                        "calibri.ttf", 
-                        "mukta-malar-extralight.ttf", 
-                        "noto-sans-arabic-medium.ttf", 
-                        "times new roman.ttf",
-                        "ANGSA.ttf",
-                        "Book-Antiqua.ttf",
-                        "CONSOLA.TTF",
-                        "COOPBL.TTF",
-                        "Rockwell-Bold.ttf",
-                        "Candara Light.TTF",
-                        "Carlito-Regular.ttf Carlito-Regular.ttf",
-                        "Castellar.ttf",
-                        "Courier New.ttf",
-                        "LSANS.TTF",
-                        "Lucida Bright Regular.ttf",
-                        "TRTempusSansITC.ttf",
-                        "Verdana.ttf",
-                        "bell-mt.ttf",
-                        "eras-itc-light.ttf",
-                        "fonnts.com-aptos-light.ttf",
-                        "georgia.ttf",
-                        "segoeuithis.ttf",
-                        "youyuan.TTF",
-                        "TfPonetoneExpanded-7BJZA.ttf",
-                    ],
-                    value="youyuan.TTF",
-                    label="Font Choice"
-                )
-                alignment = gr.Dropdown(
-                    choices=["Left", "Center", "Right", "Justified"],
-                    value="Justified",
-                    label="Text Alignment"
-                )
-                image_size = gr.Dropdown(
-                    choices=["Small", "Medium", "Large"],
-                    value="Small",
-                    label="Image Size"
-                )
-                file_format = gr.Radio(["pdf", "docx"], label="File Format", value="pdf")
-        with gr.Row():
-            get_document_btn = gr.Button(value="Get Document", elem_classes="download-btn")
-        get_document_btn.click(
-            generate_document,
-            [input_media, output_text, file_format, font_choice, font_size, line_spacing, alignment, image_size],
-            gr.File(label="Download Document")
-        )
-
-demo.launch(debug=True)
+            with gr.Row():
+                submit_btn = gr.Button("Send", variant="primary")
+                clear_btn = gr.Button("Clear", variant="secondary")
+        with gr.Column(scale=6):
+            gr.HTML('<div class="render_header"><span class="header_btn"></span><span class="header_btn"></span><span class="header_btn"></span></div>')
+            artifact_html = gr.HTML(
+                value="",
+                elem_classes="html_content"
+            )
+
+    submit_btn.click(
+        submit_chat, [chatbot, text_input], [chatbot, text_input]
+    ).then(
+        ovis_chat, [chatbot], [chatbot, artifact_html]
+    )
+
+    text_input.submit(
+        submit_chat, [chatbot, text_input], [chatbot, text_input]
+    ).then(
+        ovis_chat, [chatbot], [chatbot, artifact_html]
+    )
+
+    clear_btn.click(
+        clear_chat,
+        outputs=[chatbot, text_input, artifact_html]
+    )
+
+demo.queue(default_concurrency_limit=1).launch(server_name="0.0.0.0", share=True)