import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

# Load the tokenizer and model
model_name = "alpcansoydas/product-model-18.10.24-bert-total27label_ifhavemorethan100sampleperfamily"
tokenizer_name = "bert-base-uncased"

# Initialize tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Create a pipeline for text classification
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

# Function to classify input text
def classify_product_family(text):
    results = classifier(text)
    predicted_label = results[0]['label']
    return f"{predicted_label}"

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Product Family Classifier")
    gr.Markdown("Classify product descriptions into one of 27 family labels.")
    
    input_text = gr.Textbox(label="Enter Product Description", placeholder="Type product description here...")
    output_label = gr.Textbox(label="Predicted Family Label")
    
    classify_button = gr.Button("Classify")
    classify_button.click(fn=classify_product_family, inputs=input_text, outputs=output_label)

# Launch the Gradio interface
demo.launch()