import gradio as gr
from transformers import pipeline

# Load the model
pipe = pipeline("text-generation", model="Writer/palmyra-mini")

def chat_fn(message, history):
    # Format messages for Palmyra
    messages = [{"role": "user", "content": message}]
    
    # Generate response
    response = pipe(
        messages,
        max_new_tokens=100,
        temperature=0.7,
        do_sample=True
    )
    
    # Extract the response content
    return response[0]['generated_text']

# Messenger style interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 💬 FlzAI Messenger")
    gr.Markdown("**Created by Felix Lan**")
    
    chatbot = gr.Chatbot(
        height=500,
        show_label=False
    )
    
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Type a message...",
            show_label=False,
            container=False,
            scale=4
        )
        submit = gr.Button("Send", variant="primary", scale=1)
    
    clear = gr.Button("Clear Chat")

    def respond(message, chat_history):
        if not message.strip():
            return "", chat_history
        
        bot_message = chat_fn(message, chat_history)
        chat_history.append((message, bot_message))
        return "", chat_history

    submit.click(respond, [msg, chatbot], [msg, chatbot])
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()