import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import streamlit as st

model = AutoModelForSequenceClassification.from_pretrained('paramasivan27/RetailProductClassification_bert-base-uncased')
tokenizer = AutoTokenizer.from_pretrained('paramasivan27/RetailProductClassification_bert-base-uncased')

label_to_id = {'Electronics':0,
'Sports & Outdoors':1,
'Cell Phones & Accessories':2,
'Automotive':3,
'Toys & Games':4,
'Tools & Home Improvement':5,
'Health & Personal Care':6,
'Beauty':7,
'Grocery & Gourmet Food':8,
'Office Products':9,
'Arts, Crafts & Sewing':10,
'Pet Supplies':11,
'Patio, Lawn & Garden':12,
'Clothing, Shoes & Jewelry':13,
'Baby':14,
'Musical Instruments':15,
'Industrial & Scientific':16,
'Baby Products':17,
'Appliances':18,
'All Beauty':19,
'All Electronics':20}

id_to_label = {v: k for k, v in label_to_id.items()}

def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
    outputs = model(**inputs)
    predicted_class = torch.argmax(outputs.logits, dim=1).item()
    class_name = id_to_label[predicted_class]
    return class_name

# Streamlit app layout
st.title("Retail Product Classification")
st.write("Enter the product title and description to classify the product.")

# Input fields for title and description
title = st.text_input("Product Title")
description = st.text_area("Product Description")

# Combine title and description
if st.button("Classify Product"):
    if title and description:
        combined_text = f"{title}. {description}"
        predicted_class = predict(combined_text)
        st.write(f"Predicted Class: {predicted_class}")
    else:
        st.write("Please enter both title and description.")