import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "Luzika01883/AI"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

history = []

def chat(message):
    global history

    history.append("User: " + message)

    prompt = "\n".join(history) + "\nAssistant:"

    inputs = tokenizer(prompt, return_tensors="pt")

    output = model.generate(**inputs, max_new_tokens=120)

    reply = tokenizer.decode(output[0], skip_special_tokens=True)
    reply = reply.split("Assistant:")[-1].strip()

    history.append("Assistant: " + reply)

    return reply


with gr.Blocks() as app:
    gr.Markdown("# 💬 My AI Chatbot")

    chatbot = gr.Chatbot()
    msg = gr.Textbox()

    def respond(msg, chat_history):
        reply = chat(msg)
        chat_history.append((msg, reply))
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])

app.launch()