from transformers import pipeline
import gradio as gr

# Load the text generation pipeline
pipe = pipeline("text-generation", model="google/gemma-2-2b-jpn-it")

def generate_text(messages):
    # Extracting the content from the messages
    user_message = messages["content"]
    
    # Using the pipeline to generate a response
    result = pipe(user_message, max_length=50, num_return_sequences=1)
    
    # Return the generated text
    return result[0]["generated_text"]

# Set up Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text Generation with Hugging Face's gemma-2-2b-jpn-it")
    
    # Input for user's message
    user_input = gr.Textbox(label="Your Message", placeholder="Type a message", value="Who are you?")
    
    # Output for generated response
    output_text = gr.Textbox(label="Generated Response")
    
    # Button to trigger text generation
    generate_button = gr.Button("Generate Response")
    
    # Link button click with text generation function
    generate_button.click(fn=generate_text, inputs={"content": user_input}, outputs=output_text)

# Launch the Gradio app
demo.launch()