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
all_candidates = [
    "Soccer Jersey", "Basketball Jersey", "Football Jersey", "Baseball Jersey", "Tennis Shirt", "Hockey Jersey",
    "Soccer Ball", "Basketball", "Football", "Baseball", "Tennis Ball", "Hockey 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 Shin Guards", "Soccer Shorts"
]

def recommend(items_purchased):
    # Convert items purchased to a list and remove leading/trailing spaces
    items_purchased_list = [item.strip() for item in items_purchased.split(',')]
    
    # Filter out the purchased items from the candidates
    candidates = [item for item in all_candidates if item not in items_purchased_list]
    
    # Create the prompt
    prompt = f"ITEMS PURCHASED: {{{', '.join(items_purchased_list)}}} - CANDIDATES FOR RECOMMENDATION: {{{', '.join(candidates)}}} - RECOMMENDATION: "
    
    # Get the recommendation from the model
    model_output = t5_recommender(prompt)
    recommendation = model_output[0]['generated_text']
    return recommendation

with gr.Blocks() as demo:
    gr.Markdown("# Sports Equipment Recommender")
    
    gr.Markdown("## All Possible Candidates")
    gr.Markdown(", ".join(all_candidates))
    
    with gr.Row():
        with gr.Column():
            items_input = gr.Textbox(label="Items Purchased", placeholder="Enter a list of items purchased from the possible candidates above, separated by commas.")
        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()