| import os |
| import re |
|
|
| def main(): |
| app_path = "/home/mohammed/cbackup/Coding/R&D/SaasBackend/New/TextExtractor-v1/app.py" |
| |
| with open(app_path, "r") as f: |
| original_code = f.read() |
| |
| commented_backup = "\n".join("# " + line for line in original_code.splitlines()) |
| |
| new_code = """# Main One |
| import json |
| import os |
| import re |
| from datetime import datetime |
| |
| import fitz # PyMuPDF |
| import gradio as gr |
| import spaces |
| import torch |
| from gradio.themes.base import Base |
| from PIL import Image |
| from qwen_vl_utils import process_vision_info |
| from transformers import AutoProcessor, Qwen2VLForConditionalGeneration |
| |
| |
| # 1. Custom Theme Definition |
| class CustomTheme(Base): |
| def __init__(self): |
| super().__init__() |
| self.primary_hue = "blue" |
| self.secondary_hue = "sky" |
| |
| |
| custom_theme = CustomTheme() |
| |
| DESCRIPTION = "A powerful vision-language model that can understand images and text to provide detailed analysis." |
| |
| |
| # 2. Safely Downscale & Save Image to prevent CUDA OOM |
| def prepare_and_save_image(image_filepath, max_width=1250, max_height=1750): |
| if not image_filepath or not os.path.exists(image_filepath): |
| raise ValueError("Image file not found.") |
| |
| img = Image.open(image_filepath).convert("RGB") |
| width, height = img.size |
| |
| # Re-calculate dimensions while locking aspect ratio |
| if width > max_width or height > max_height: |
| aspect_ratio = width / height |
| if width > max_width: |
| width = max_width |
| height = int(width / aspect_ratio) |
| if height > max_height: |
| height = max_height |
| width = int(height * aspect_ratio) |
| |
| img = img.resize((width, height), Image.Resampling.LANCZOS) |
| |
| # We MUST save the resized image to a new path so the GPU actually reads the small version |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") |
| temp_filename = os.path.abspath(f"temp_downscaled_{timestamp}.png") |
| img.save(temp_filename, "PNG") |
| |
| return temp_filename, width, height |
| |
| |
| # 3. PDF Page Extractor |
| def convert_pdf_to_images(pdf_path): |
| image_paths = [] |
| doc = fitz.open(pdf_path) |
| base_name = os.path.splitext(os.path.basename(pdf_path))[0] |
| |
| for i, page in enumerate(doc): |
| # dpi=150 is the sweet spot for 7B models to read fine text without blowing out VRAM |
| pix = page.get_pixmap(dpi=150) |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| image_path = os.path.abspath(f"{base_name}_page_{i + 1}_{timestamp}.png") |
| pix.save(image_path) |
| image_paths.append(image_path) |
| |
| doc.close() |
| return image_paths |
| |
| |
| # 4. Bulletproof JSON Extractor |
| def extract_json_from_text(raw_text): |
| # Target 1: Look inside markdown json fences |
| match = re.search(r"\\`\\`\\`(?:json)?\\s*(\\{.*?\\})\\s*\\`\\`\\`", raw_text, re.DOTALL) |
| if match: |
| try: |
| return json.loads(match.group(1)) |
| except json.JSONDecodeError: |
| pass |
| |
| # Target 2: Fallback to raw bracket math |
| try: |
| start = raw_text.find("{") |
| end = raw_text.rfind("}") + 1 |
| if start != -1 and end > start: |
| return json.loads(raw_text[start:end]) |
| except json.JSONDecodeError: |
| pass |
| |
| return None |
| |
| def extract_html_from_text(raw_text): |
| match = re.search(r"\\`\\`\\`(?:html)?\\s*(<html.*?>.*?</html>)\\s*\\`\\`\\`", raw_text, re.DOTALL | re.IGNORECASE) |
| if match: |
| return match.group(1) |
| |
| if "<html" in raw_text.lower(): |
| start = raw_text.lower().find("<html") |
| end = raw_text.lower().rfind("</html>") + 7 |
| if start != -1 and end > start: |
| return raw_text[start:end] |
| |
| return raw_text |
| |
| |
| # 5. Global Model Init (Optimized with SDPA & bfloat16) |
| model = Qwen2VLForConditionalGeneration.from_pretrained( |
| "Qwen/Qwen2-VL-7B-Instruct", torch_dtype=torch.bfloat16, attn_implementation="sdpa" |
| ) |
| processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") |
| |
| |
| @spaces.GPU(duration=180) |
| def run_inference(uploaded_files, text_input): |
| if not uploaded_files: |
| err = json.dumps({"error": "No file uploaded."}, indent=4) |
| return err, gr.Button(interactive=False) |
| |
| results = [] |
| files_to_delete_from_disk = [] |
| |
| # Standardize incoming Gradio file objects to raw string paths |
| raw_paths = [getattr(f, "path", getattr(f, "name", str(f))) for f in uploaded_files] |
| |
| images_to_process = [] |
| unsupported = [] |
| |
| # Sort files into PDFs vs standard images |
| for f_path in raw_paths: |
| ext = os.path.splitext(f_path)[1].lower() |
| if ext == ".pdf": |
| try: |
| generated_pngs = convert_pdf_to_images(f_path) |
| images_to_process.extend(generated_pngs) |
| files_to_delete_from_disk.extend(generated_pngs) |
| except Exception as e: |
| results.append( |
| json.dumps( |
| {"error": f"Corrupt PDF: {os.path.basename(f_path)}"}, |
| indent=4, |
| ) |
| ) |
| elif ext in [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"]: |
| images_to_process.append(f_path) |
| else: |
| unsupported.append(os.path.basename(f_path)) |
| |
| if unsupported: |
| results.append( |
| json.dumps( |
| {"warning": f"Ignored unknown files: {', '.join(unsupported)}"}, |
| indent=4, |
| ) |
| ) |
| |
| system_json_injection = ( |
| f"{text_input}\\n\\nBased on the image and the query, respond ONLY with a single, " |
| "valid JSON object. This object should be well-structured, using nested objects " |
| "and arrays to logically represent the information." |
| ) |
| |
| for original_img in images_to_process: |
| downscaled_img = None |
| try: |
| downscaled_img, w, h = prepare_and_save_image(original_img) |
| files_to_delete_from_disk.append(downscaled_img) |
| |
| messages = [ |
| { |
| "role": "user", |
| "content": [ |
| { |
| "type": "image", |
| "image": downscaled_img, |
| "resized_height": h, |
| "resized_width": w, |
| }, |
| {"type": "text", "text": system_json_injection}, |
| ], |
| } |
| ] |
| |
| text = processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True |
| ) |
| image_inputs, video_inputs = process_vision_info(messages) |
| |
| inputs = processor( |
| text=[text], |
| images=image_inputs, |
| videos=video_inputs, |
| padding=True, |
| return_tensors="pt", |
| ).to("cuda") |
| |
| # Optimized generation parameters (Faster + Better JSON) |
| generated_ids = model.generate( |
| **inputs, max_new_tokens=2048, do_sample=False, use_cache=True |
| ) |
| trimmed = [ |
| out[len(in_ids) :] |
| for in_ids, out in zip(inputs.input_ids, generated_ids) |
| ] |
| raw_output = processor.batch_decode( |
| trimmed, |
| skip_special_tokens=True, |
| clean_up_tokenization_spaces=True, |
| )[0] |
| |
| # Format clean output |
| parsed_json = extract_json_from_text(raw_output) |
| clean_source_name = re.sub( |
| r"_\\d{8}_\\d{6}\\.png$", "", os.path.basename(original_img) |
| ) |
| |
| if parsed_json: |
| parsed_json["_source_document"] = clean_source_name |
| results.append(json.dumps(parsed_json, indent=4)) |
| else: |
| results.append( |
| json.dumps( |
| { |
| "error": "Model failed to format valid JSON", |
| "source": clean_source_name, |
| "raw_text": raw_output[:250] + "...", |
| }, |
| indent=4, |
| ) |
| ) |
| |
| except Exception as e: |
| results.append( |
| json.dumps( |
| { |
| "error": f"Inference failed on {os.path.basename(original_img)}", |
| "trace": str(e), |
| }, |
| indent=4, |
| ) |
| ) |
| |
| # Rigorous disk sweep: Delete all generated temp files |
| for filepath in set(files_to_delete_from_disk): |
| if filepath and os.path.exists(filepath): |
| try: |
| os.remove(filepath) |
| except OSError: |
| pass |
| |
| final_payload = "\\n\\n".join(results) |
| is_failed = '"error":' in final_payload |
| |
| return final_payload, gr.Button(interactive=not is_failed) |
| |
| |
| @spaces.GPU(duration=180) |
| def run_html_replica(uploaded_files): |
| if not uploaded_files: |
| err = "<!-- Error: No file uploaded. -->" |
| return err, err |
| |
| files_to_delete_from_disk = [] |
| |
| # Standardize incoming Gradio file objects to raw string paths |
| raw_paths = [getattr(f, "path", getattr(f, "name", str(f))) for f in uploaded_files] |
| |
| images_to_process = [] |
| |
| # Sort files into PDFs vs standard images |
| for f_path in raw_paths: |
| ext = os.path.splitext(f_path)[1].lower() |
| if ext == ".pdf": |
| try: |
| generated_pngs = convert_pdf_to_images(f_path) |
| images_to_process.extend(generated_pngs) |
| files_to_delete_from_disk.extend(generated_pngs) |
| except Exception as e: |
| pass |
| elif ext in [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"]: |
| images_to_process.append(f_path) |
| |
| if not images_to_process: |
| err = "<!-- Error: No valid image or PDF found. -->" |
| return err, err |
| |
| system_html_injection = ( |
| "You are an expert frontend web developer. Your task is to recreate the provided image exactly as a single HTML file containing inline CSS. " |
| "Replicate the color, font, theme, alignment, and icons perfectly (1:1 replica). " |
| "Output ONLY valid HTML code starting with <html>. Do not include markdown formatting like ```html." |
| ) |
| |
| final_html_parts = [] |
| |
| for original_img in images_to_process: |
| downscaled_img = None |
| try: |
| downscaled_img, w, h = prepare_and_save_image(original_img) |
| files_to_delete_from_disk.append(downscaled_img) |
| |
| messages = [ |
| { |
| "role": "user", |
| "content": [ |
| { |
| "type": "image", |
| "image": downscaled_img, |
| "resized_height": h, |
| "resized_width": w, |
| }, |
| {"type": "text", "text": system_html_injection}, |
| ], |
| } |
| ] |
| |
| text = processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True |
| ) |
| image_inputs, video_inputs = process_vision_info(messages) |
| |
| inputs = processor( |
| text=[text], |
| images=image_inputs, |
| videos=video_inputs, |
| padding=True, |
| return_tensors="pt", |
| ).to("cuda") |
| |
| generated_ids = model.generate( |
| **inputs, max_new_tokens=4096, do_sample=False, use_cache=True |
| ) |
| trimmed = [ |
| out[len(in_ids) :] |
| for in_ids, out in zip(inputs.input_ids, generated_ids) |
| ] |
| raw_output = processor.batch_decode( |
| trimmed, |
| skip_special_tokens=True, |
| clean_up_tokenization_spaces=True, |
| )[0] |
| |
| parsed_html = extract_html_from_text(raw_output) |
| final_html_parts.append(parsed_html) |
| |
| except Exception as e: |
| final_html_parts.append(f"<!-- Inference failed on {os.path.basename(original_img)}: {str(e)} -->") |
| |
| # Rigorous disk sweep: Delete all generated temp files |
| for filepath in set(files_to_delete_from_disk): |
| if filepath and os.path.exists(filepath): |
| try: |
| os.remove(filepath) |
| except OSError: |
| pass |
| |
| final_payload = "\\n<hr/>\\n".join(final_html_parts) |
| return final_payload, final_payload |
| |
| |
| @spaces.GPU(duration=180) |
| def generate_explanation(json_text): |
| if not json_text or '"error":' in json_text: |
| return "Cannot generate an explanation from an errored JSON payload." |
| |
| prompt = ( |
| "You are an expert data analyst. Your task is to provide a comprehensive, human-readable explanation " |
| "of the following JSON data, which may represent one or more pages from a document. First, provide a textual explanation. " |
| "so the json which is provided try to understand what it is representing like a receipt, table, or list of items. or just some text or just an image and after getting the context then only provide the explanation." |
| "If the JSON contains data from multiple sources (pages), explain each one. Then, if the JSON data represents a table, " |
| "a list of items, or a receipt, you **must** re-format the key information into a Markdown table for clarity.\\n\\n" |
| f"JSON Data:\\n```json\\n{json_text}\\n```" |
| ) |
| |
| messages = [{"role": "user", "content": prompt}] |
| text = processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True |
| ) |
| inputs = processor(text=[text], return_tensors="pt").to("cuda") |
| |
| generated_ids = model.generate( |
| **inputs, max_new_tokens=1536, do_sample=False, use_cache=True |
| ) |
| trimmed = [ |
| out[len(in_ids) :] for in_ids, out in zip(inputs.input_ids, generated_ids) |
| ] |
| return processor.batch_decode(trimmed, skip_special_tokens=True)[0] |
| |
| |
| # 6. Gradio UI Assembly |
| css = \"\"\" |
| .gradio-container { font-family: 'IBM Plex Sans', sans-serif; } |
| |
| #output-code, #output-code pre, #output-code code { |
| background-color: #f0f0f0; |
| border: 1px solid #e0e0e0; |
| border-radius: 7px; |
| color: #333; |
| } |
| #output-code .token.punctuation { color: #393a34; } |
| #output-code .token.property, #output-code .token.string { color: #0b7500; } |
| #output-code .token.number { color: #2973b7; } |
| #output-code .token.boolean { color: #9a050f; } |
| |
| #explanation-box { |
| min-height: 200px; |
| border: 1px solid #e0e0e0; |
| padding: 15px; |
| border-radius: 7px; |
| } |
| |
| .dark #output-code, .dark #output-code pre, .dark #output-code code { |
| background-color: #2b2b2b !important; |
| border: 1px solid #444 !important; |
| color: #f0f0f0 !important; |
| } |
| .dark #explanation-box { border: 1px solid #444 !important; } |
| .dark #output-code code span { color: #f0f0f0 !important; } |
| .dark #output-code .token.punctuation { color: #ccc !important; } |
| .dark #output-code .token.property, .dark #output-code .token.string { color: #90ee90 !important; } |
| .dark #output-code .token.number { color: #add8e6 !important; } |
| .dark #output-code .token.boolean { color: #f08080 !important; } |
| \"\"\" |
| |
| with gr.Blocks(theme=custom_theme, css=css) as demo: |
| gr.Markdown("# Sparrow Qwen2-VL-7B Vision AI ๐๏ธ") |
| gr.Markdown(DESCRIPTION) |
| |
| with gr.Tabs(): |
| with gr.Tab("JSON Extraction"): |
| with gr.Row(): |
| with gr.Column(scale=1): |
| input_files = gr.Files( |
| label="Upload Images or PDFs", |
| file_types=[ |
| ".pdf", |
| ".png", |
| ".jpg", |
| ".jpeg", |
| ".bmp", |
| ".gif", |
| ".webp", |
| ], |
| ) |
| text_input = gr.Textbox( |
| label="Your Query", |
| placeholder="e.g., Extract all line items into JSON.", |
| ) |
| submit_btn = gr.Button("Analyze File(s)", variant="primary") |
| |
| with gr.Column(scale=2): |
| output_text = gr.Code( |
| label="Full JSON Response", |
| language="json", |
| elem_id="output-code", |
| interactive=False, |
| ) |
| explanation_btn = gr.Button( |
| "๐ Generate Detailed Explanation", interactive=False |
| ) |
| explanation_output = gr.Markdown( |
| label="Detailed Explanation", elem_id="explanation-box" |
| ) |
| |
| submit_btn.click( |
| fn=run_inference, |
| inputs=[input_files, text_input], |
| outputs=[output_text, explanation_btn], |
| api_name="analyze_document", |
| ) |
| |
| explanation_btn.click( |
| fn=generate_explanation, |
| inputs=[output_text], |
| outputs=[explanation_output], |
| api_name="generate_explanation", |
| ) |
| |
| with gr.Tab("HTML Replica"): |
| with gr.Row(): |
| with gr.Column(scale=1): |
| html_input_files = gr.Files( |
| label="Upload Images or PDFs", |
| file_types=[ |
| ".pdf", |
| ".png", |
| ".jpg", |
| ".jpeg", |
| ".bmp", |
| ".gif", |
| ".webp", |
| ], |
| ) |
| html_submit_btn = gr.Button("Generate HTML Replica", variant="primary") |
| |
| with gr.Column(scale=2): |
| html_rendered_output = gr.HTML( |
| label="Rendered HTML Replica" |
| ) |
| html_raw_output = gr.Code( |
| label="Raw HTML Source", |
| language="html", |
| interactive=False, |
| ) |
| |
| html_submit_btn.click( |
| fn=run_html_replica, |
| inputs=[html_input_files], |
| outputs=[html_rendered_output, html_raw_output], |
| api_name="generate_html_replica", |
| ) |
| |
| if __name__ == "__main__": |
| demo.queue(api_open=True).launch(debug=True) |
| """ |
| |
| with open(app_path, "w") as f: |
| f.write(commented_backup) |
| f.write("\n\n") |
| f.write(new_code) |
|
|
| if __name__ == "__main__": |
| main() |
|
|