RedaAlami's picture
Update app.py
608c6fe verified
raw
history blame
2.12 kB
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")
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()