import gradio as gr
from transformers import pipeline

# Load a text generation pipeline from Hugging Face
rap_generator = pipeline("text-generation", model="gpt2")  # You can use a rap-specific model if available

def generate_rap(quantity, style, countries, theme, word_list, freestyle):
    # Format the prompt for the AI model
    prompt = (
        f"{quantity} lines of {style} rap inspired by {' and '.join(countries)} "
        f"about {theme}. Featuring words like {' '.join(word_list)}. "
        f"{'Freestyle' if freestyle else 'Structured'}."
    )
    # Generate rap lyrics using the model
    response = rap_generator(prompt, max_length=50, num_return_sequences=1)
    return response[0]['generated_text']

demo = gr.Interface(
    generate_rap,
    [
        gr.Slider(2, 20, value=4, label="Number of Lines", info="Choose the number of lines for your rap."),
        gr.Dropdown(
            ["chill", "hardcore", "freestyle", "battle"], label="Rap Style", info="Select the style of rap."
        ),
        gr.CheckboxGroup(["USA", "Japan", "UK", "France"], label="Countries", info="Where's the rap vibe from?"),
        gr.Radio(["love", "struggle", "success", "hustle"], label="Theme", info="What’s the rap about?"),
        gr.Dropdown(
            ["money", "dream", "fight", "night"], value=["money", "fight"], multiselect=True,
            label="Key Words", info="Words to include in the rap."
        ),
        gr.Checkbox(label="Freestyle", info="Make it a freestyle?")
    ],
    "text",
    examples=[
        [4, "chill", ["USA"], "love", ["money", "dream"], True],
        [10, "hardcore", ["UK", "Japan"], "hustle", ["fight", "night"], False],
        [8, "battle", ["France"], "struggle", ["dream"], True],
    ]
)

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