import gradio as gr
from transformers import pipeline

# Load fine-tuned model from Hugging Face Hub
t5_recommender = pipeline(model="RedaAlami/t5_recommendation_sports_equipment_english")

# Fixed list of candidates
candidates = (
    "Soccer Jersey, Basketball Jersey, Football Jersey, Baseball Jersey, Tennis Shirt, "
    "Hockey Jersey, Soccer Ball, Basketball, Football, Baseball, Tennis Ball, Hocket Puck, "
    "Soccer Cleats, Basketball Shoes, Football Cleats, Baseball Cleats, Tennis Shoes, Hockey Helmet, "
    "Goalie Gloves, Basketball Arm Sleeve, Football Shoulder Pads, Baseball Cap, Tennis Racket, Hockey Skates, "
    "Soccer Goal Post, Basketball Hoop, Football Helmet, Baseball Bat, Hockey Stick, Soccer Cones, Basketball Shorts, "
    "Baseball Glove, Hockey Pads, Soccer Shorts"
)

def recommend(items_purchased):
    prompt = f"ITEMS PURCHASED: {{{items_purchased}}} - CANDIDATES FOR RECOMMENDATION: {{{candidates}}} - RECOMMENDATION: "
    model_output = t5_recommender(prompt)
    recommendation = model_output[0]['generated_text']
    return recommendation

with gr.Blocks() as demo:
    gr.Markdown("# Sports Equipment Recommender")
    with gr.Row():
        with gr.Column():
            items_input = gr.Textbox(label="Items Purchased")
        with gr.Column():
            recommendation_output = gr.Textbox(label="Recommendation")

    recommend_button = gr.Button("Get Recommendation")
    recommend_button.click(fn=recommend, inputs=items_input, outputs=recommendation_output)

demo.launch()