import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import spaces

model_name = "MBZUAI-Paris/Atlas-Chat-9B"
dtype = torch.bfloat16
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
     torch_dtype=dtype,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)


@spaces.GPU
def chat(input_text, history=[]):
    # Tokenize the input and generate response
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=150)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Update the conversation history
    history.append((input_text, response))
    return history, history


with gr.Blocks() as app:
    gr.Markdown("<h1 style='text-align: center;'>دردشة أطلس</h1>")
    chatbot = gr.Chatbot(label="المحادثة")
    state = gr.State([])

    with gr.Row():
        txt = gr.Textbox(show_label=False, placeholder="اكتب رسالتك هنا...").style(container=False)
        send_button = gr.Button("إرسال")

    # Define the button click event
    def user_input(input_text, history):
        return "", history + [[input_text, None]]

    def bot_response(history):
        user_message = history[-1][0]
        inputs = tokenizer(user_message, return_tensors="pt").to(model.device)
        outputs = model.generate(**inputs, max_new_tokens=150)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        history[-1][1] = response
        return history

    # Link functions to button clicks
    txt.submit(user_input, [txt, state], [txt, state]).then(
        bot_response, state, chatbot
    )
    send_button.click(user_input, [txt, state], [txt, state]).then(
        bot_response, state, chatbot
    )

app.launch()