import gradio as gr
from transformers import pipeline
from PIL import Image
import pytesseract

# Initialize chat model
chat_model = pipeline("text-generation", model="gpt2")  # عدّل اسم النموذج حسب الحاجة

# Chat function
def chat_fn(history, user_input):
    conversation = {"history": history, "user": user_input}
    response = chat_model(user_input, max_length=50, num_return_sequences=1)
    conversation["bot"] = response[0]['generated_text']
    history.append((user_input, conversation["bot"]))
    return history, ""

# OCR function
def ocr(image):
    text = pytesseract.image_to_string(image)
    return text

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("### الصور والدردشة")

    # Image OCR section
    with gr.Tab("استخراج النصوص من الصور"):
        with gr.Row():
            image_input = gr.Image(type="pil")
            ocr_output = gr.Textbox()
        submit_button = gr.Button("Submit")
        submit_button.click(ocr, inputs=image_input, outputs=ocr_output)

    # Chat section
    with gr.Tab("المحادثة"):
        chatbot = gr.Chatbot()
        msg = gr.Textbox(label="اكتب رسالتك")
        clear = gr.Button("Clear")
        msg.submit(chat_fn, [chatbot, msg], [chatbot, msg])
        clear.click(lambda: None, None, chatbot)

demo.launch()