Update space

app.py

@@ -6,11 +6,36 @@ 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()
-    return predicted_class
+    class_name = id_to_label[predicted_class]
+    return class_name
 
 # Streamlit app layout
 st.title("Retail Product Classification")