import gradio as gr
from transformers import pipeline

# Load the text generation model
generator = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", device_map="auto", max_new_tokens=200)

def respond(message, chat_history):
    prompt = f"User: {message}\nAssistant:"
    result = generator(prompt)[0]["generated_text"]
    response = result.split("Assistant:")[-1].strip()
    chat_history.append((message, response))
    return "", chat_history

# Interface: PromptVerse
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    with gr.Row():
        msg = gr.Textbox(placeholder="Type your message...", show_label=False, scale=8)
        submit = gr.Button(value="Submit", scale=2)
    state = gr.State([])

    submit.click(fn=respond, inputs=[msg, state], outputs=[msg, chatbot])
    msg.submit(fn=respond, inputs=[msg, state], outputs=[msg, chatbot])

demo.launch()