app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoProcessor, AutoModelForImageTextToText
from PIL import Image
from pdf2image import convert_from_bytes
import torch
from io import BytesIO

# Model bilgisi
model_name = "ynsbyrm/Nanonets-OCR2-3B-demo-clone"

# Model ve processor yükleme
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.truncation_side = "left"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForImageTextToText.from_pretrained(
    model_name,
    dtype=torch.float16 if device=="cuda" else torch.float32,
    device_map="auto",
    trust_remote_code=True
)
model.eval()
processor = AutoProcessor.from_pretrained(model_name)

def build_prompt(prompt_text):
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": prompt_text}
            ]
        }
    ]
    prompt = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    return prompt

# OCR fonksiyonu
def ocr_file(file):
    prompt_text = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
    prompt = build_prompt(prompt_text)
    texts = []
    print("Filename:" + file.name)
    if file.name.endswith(".pdf"):
        print("in if")
        # PDF -> sayfalara ayır
        pages = convert_from_bytes(file)
        for i, page in enumerate(pages):
            inputs = processor(text=[prompt], images=[page], return_tensors="pt").to(device)
            output_ids = model.generate(**inputs, max_new_tokens=1024)
            text = processor.tokenizer.decode(output_ids[0], skip_special_tokens=True)
            texts.append(f"--- Sayfa {i+1} ---\n{text}")
    else:
        print("in else")
        # PNG / JPG
        img = Image.open(file).convert("RGB")
        print("image opened")
        inputs = processor(
    		text=[prompt],
            images=[img],
            return_tensors="pt",
            padding=True
        ).to(model.device)
        print("inputs:", inputs)
    
        output_ids = model.generate(
    		**inputs,
    		max_new_tokens=1024, 
    		do_sample=False
    	)
        print("output_ids generated")

        output_text = processor.batch_decode(
            output_ids,
            skip_special_tokens=True
        )
        print("--------------------\n\n\nbatch decode:", output_text)
        print("\n\n\n-------------")
        
        #output_text = processor.batch_decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
        return output_text[0]

# Gradio arayüzü
gr.Interface(
    fn=ocr_file,
    inputs=gr.File(file_types=[".pdf", ".png", ".jpg"]),
    outputs=gr.Textbox(label="OCR Sonucu", lines=20),
    title="Model Tabanlı OCR Arayüzü",
    description="PDF veya resim yükleyin, model tabanlı OCR sonucu alın."
).launch()