added app.py

app.py

@@ -0,0 +1,80 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+from PIL import Image
+import requests
+import gradio as gr
+import pandas as pd
+import subprocess
+import os
+
+# Install flash-attn without CUDA build
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+# Load the model and processor
+model_id = "yifeihu/TB-OCR-preview-0.1"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    device_map="cuda", 
+    trust_remote_code=True, 
+    torch_dtype="auto", 
+    attn_implementation='flash_attention_2',
+    load_in_4bit=True
+)
+processor = AutoProcessor.from_pretrained(model_id, 
+    trust_remote_code=True, 
+    num_crops=16
+)
+
+# Define the OCR function
+def phi_ocr(image):
+    question = "Convert the text to markdown format."
+    prompt_message = [{
+        'role': 'user',
+        'content': f'<|image_1|>\n{question}',
+    }]
+    prompt = processor.tokenizer.apply_chat_template(prompt_message, tokenize=False, add_generation_prompt=True)
+    inputs = processor(prompt, [image], return_tensors="pt").to("cuda")
+    generation_args = { 
+        "max_new_tokens": 1024, 
+        "temperature": 0.1, 
+        "do_sample": False
+    }
+    generate_ids = model.generate(**inputs, eos_token_id=processor.tokenizer.eos_token_id, **generation_args)
+    generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
+    response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    response = response.split("<image_end>")[0]
+    return response
+
+# Define the function to process multiple images and save results to a CSV
+def process_images(input_images):
+    results = []
+    for index, image in enumerate(input_images):
+        extracted_text = phi_ocr(image)
+        results.append({
+            'index': index,
+            'extracted_text': extracted_text
+        })
+    
+    # Convert to DataFrame and save to CSV
+    df = pd.DataFrame(results)
+    output_csv = "extracted_entities.csv"
+    df.to_csv(output_csv, index=False)
+    
+    return f"Processed {len(input_images)} images and saved to {output_csv}", output_csv
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# OCR with TB-OCR-preview-0.1")
+    gr.Markdown("Upload multiple images to extract and convert text to markdown format.")
+    gr.Markdown("[Check out the model here](https://huggingface.co/yifeihu/TB-OCR-preview-0.1)")
+    
+    with gr.Row():
+        input_images = gr.Image(type="pil", label="Upload Images", tool="editor", source="upload", multiple=True)
+        output_text = gr.Textbox(label="Status")
+        output_csv_link = gr.File(label="Download CSV")
+
+    input_images.change(fn=process_images, inputs=input_images, outputs=[output_text, output_csv_link])
+
+demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+transformers
+torch
+torchvision
+Pillow
+pandas
+gradio
+flash-attn