import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import pytesseract

# Use a pipeline as a high-level helper
from transformers import pipeline

# Initialize tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("sambanovasystems/SambaLingo-Arabic-Chat")
model = AutoModelForCausalLM.from_pretrained("sambanovasystems/SambaLingo-Arabic-Chat")

# Chat function
def chat_fn(history, user_input):
    conversation = {"history": history, "user": user_input}
    # Generate a response using the model
    input_ids = tokenizer.encode(user_input, return_tensors="pt")
    response = model.generate(input_ids=input_ids, max_length=50)
    conversation["bot"] = tokenizer.decode(response[0], skip_special_tokens=True)
    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)

# Launch the Gradio interface
demo.launch()