import os
from dotenv import load_dotenv
from huggingface_hub import InferenceClient
import gradio as gr
from PIL import Image
from io import BytesIO

load_dotenv()

# Initialize Hugging Face Inference client
hf_token = os.getenv("HF_API_TOKEN")
client = InferenceClient(token=hf_token)

def run_ocr(file):
    if file is None:
        return "No file uploaded"
    
    # Read uploaded file
    data = file.read()
    
    # Call the Typhoon OCR model
    result = client.text_generation(
        model="scb10x/typhoon-ocr1.5-2b",
        inputs=data
    )
    
    # Extract generated text
    ocr_text = result[0]["generated_text"]
    
    return ocr_text

# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Typhoon OCR Web App")
    file_input = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg"])
    output_text = gr.Textbox(label="OCR Result", lines=20)
    run_btn = gr.Button("Run OCR")
    
    run_btn.click(fn=run_ocr, inputs=file_input, outputs=output_text)

demo.launch()