Stanford-TH's picture
Update app.py
19b9239 verified
raw
history blame
682 Bytes
import gradio as gr
from transformers import pipeline
# Load the genre prediction model as a pipeline
pipe = pipeline(model="Stanford-TH/GenrePrediction", trust_remote_code=True)
def classify_movie_genre(description):
# Get predictions using the pipeline
predictions = pipe(description)
# Format the predictions to match Gradio's output expectations
formatted_predictions = {genre: score for genre, score in predictions}
return formatted_predictions
# Define the Gradio interface
iface = gr.Interface(
fn=classify_movie_genre,
inputs="text",
live=True,
title="Genre Prediction"
)
# Launch the Gradio interface
iface.launch(inline=False)