import gradio as gr 
from transformers import pipeline

task = "text-generation"  # Specify the task as text generation
model = "tiiuae/falcon-7b-instruct"  # Use a specific model (option 1) or
# model = "tiiuae/falcon-40b-instruct"  # Use another specific model (option 2)
get_completion = pipeline(task=task, model=model, trust_remote_code=True)  # Create a text generation pipeline

# device = "cuda" if torch.cuda.is_available() else "cpu"

# Define a function to format the chat history and create a prompt for the chatbot
def format_chat_prompt(message, chat_history):
    prompt = ""
    for turn in chat_history:
        user_message, bot_message = turn
        prompt = f"{prompt}\nUser: {user_message}\nAssistant: {bot_message}"
    prompt = f"{prompt}\nUser: {message}\nAssistant:"
    return prompt

# Define a function to respond to the user's message
def respond(message, chat_history):
    formatted_prompt = format_chat_prompt(message, chat_history)
    # Generate a response using the specified model, given the formatted prompt
    bot_message = get_completion(formatted_prompt,
                                 max_new_tokens=1024,
                                 # stop_sequences=["\nUser:", "<|endoftext|>"],
                                 stop_sequences=["\nUser:", ""]).generated_text
    chat_history.append((message, bot_message))  # Add the user and bot messages to the chat history
    return "", chat_history

# Create a Gradio interface with UI components
with gr.Blocks() as demo:
    chatbot = gr.Chatbot(height=240)  # Create a chatbot UI component
    msg = gr.Textbox(label="Prompt")  # Create a textbox UI component for user input
    btn = gr.Button("Submit")  # Create a submit button UI component
    # Create a clear button UI component that clears the console
    clear = gr.ClearButton(components=[msg, chatbot], value="Clear console")

    # Define the behavior of the submit button when clicked
    btn.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
    # Define the behavior of the enter key in the textbox (submit action)
    msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])

# Launch the Gradio interface
demo.launch()