File size: 1,779 Bytes
6f69f46 7b0a595 6f69f46 7b0a595 6f69f46 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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.")
|