import gradio as gr
from transformers import pipeline

# Load the model
model_id = "LiquidAI/LFM2.5-350M"
pipe = pipeline("text-generation", model=model_id, device_map="auto")

def simple_chat(user_input):
    # Simplified prompt without character/scenario
    prompt = f"### User: {user_input}\n### Assistant:"
    
    response = pipe(prompt, max_new_tokens=150, temperature=0.7, top_p=0.9, do_sample=True)
    # Clean the output to show only the assistant's reply
    return response[0]['generated_text'].split("### Assistant:")[-1].strip()

# Create a clean, simple UI
with gr.Blocks() as demo:
    gr.Markdown("# 💬 Simple AI Chat")
    
    input_text = gr.Textbox(label="Your Message")
    output_text = gr.Textbox(label="AI Response", interactive=False)
    
    btn = gr.Button("Send")
    btn.click(fn=simple_chat, inputs=input_text, outputs=output_text)

demo.launch()