app.py

import gradio as gr
import torch
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

# Model and tokenizer config
MODEL_REPO_ID = "LGAI-EXAONE/EXAONE-4.0-1.2B"

print("✅ Starting application...")

try:
    print(f"🔄 Loading tokenizer from '{MODEL_REPO_ID}'...")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_REPO_ID,
        trust_remote_code=True  
    )
    print("✅ Tokenizer loaded successfully.")

    print(f"🔄 Loading model '{MODEL_REPO_ID}' with torch_dtype=torch.bfloat16...")

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_REPO_ID,
        torch_dtype=torch.bfloat16,
        device_map="auto", 
        trust_remote_code=True  
    )
    print("✅ Model loaded successfully.")

except Exception as e:
    print(f"❌ Error loading model or tokenizer: {e}")
    raise

# Streaming chat function 
def user_input_handler(user_message, history):
    return "", history + [[user_message, None]]

def bot_stream(history):
    user_message = history[-1][0]
    history[-1][1] = "" # Initialize the bot's response field

    # Format the conversation history into the model's expected chat format
    messages = []
    for human, assistant in history[:-1]:
        messages.append({"role": "user", "content": human})
        if assistant:
            messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": user_message})

    # Apply the chat template to format the prompt correctly
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    
    # Tokenize the formatted prompt
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # Use TextIteratorStreamer for non-blocking, token-by-token generation
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    # Set up the generation parameters in a dictionary
    generation_kwargs = dict(
        inputs,
        streamer=streamer,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id,
    )

    # Run the generation in a separate thread
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # Yield each new token to the Gradio chat interface as it's generated
    for token in streamer:
        history[-1][1] += token
        yield history

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), css="footer {display: none !important}") as demo:
    gr.Markdown("## 🤖 EXAONE-4.0-1.2B")
    gr.Markdown("This demo runs `LGAI-EXAONE/EXAONE-4.0-1.2B` model")

    chatbot = gr.Chatbot(label="Chat History", height=600, bubble_full_width=False)
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Type your message here...",
            label="Your Message",
            scale=8,
            autofocus=True,
        )
        send_btn = gr.Button("Send", scale=1, variant="primary")

    clear_btn = gr.ClearButton([msg, chatbot], value="🗑️ Clear Chat")

    # Event Handlers - when a message is submitted 
    msg.submit(user_input_handler, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_stream, chatbot, chatbot
    )
    send_btn.click(user_input_handler, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_stream, chatbot, chatbot
    )

demo.queue()
demo.launch(debug=True)