Update app.py

app.py

@@ -10,10 +10,10 @@ import nltk
 nltk.download('stopwords')
 from nltk.corpus import stopwords
 
-# Load artifacts
-model_path = "./best_product_model"
+# Setup
+model_path = "."  # All files are in the root directory
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-product_encoder = joblib.load(f"{model_path}/product_encoder.pkl")
+product_encoder = joblib.load("category_encoder.pkl")
 base_model_name = "DataScienceWFSR/bert-food-product-category-cw"
 stop_words = set(stopwords.words('english'))
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -32,7 +32,7 @@ def clean_text(text):
 def template_(day, month, year, country, title, text):
     return f"Date: day {day}, month {month}, year {year}. Country: {country}. Title: {title}. Text: {text}"
 
-# Model class
+# Model definition
 class ProductCategoryClassifier(nn.Module):
     def __init__(self, model_name, num_categories):
         super().__init__()
@@ -50,7 +50,7 @@ class ProductCategoryClassifier(nn.Module):
 # Load model
 num_categories = len(product_encoder.classes_)
 model = ProductCategoryClassifier(model_name=base_model_name, num_categories=num_categories).to(device)
-model.load_state_dict(torch.load(f"{model_path}/pytorch_model.bin", map_location=device))
+model.load_state_dict(torch.load("pytorch_model.bin", map_location=device))
 model.eval()
 
 # Inference function
@@ -58,19 +58,19 @@ def predict_category(day, month, year, title, text, country="Unknown"):
     title_clean = clean_text(title)
     text_clean = clean_text(text)
     input_text = template_(day, month, year, country, title_clean, text_clean)
-    
+
     inputs = tokenizer([input_text], padding=True, truncation=True, max_length=512, return_tensors="pt")
     input_ids = inputs['input_ids'].to(device)
     attention_mask = inputs['attention_mask'].to(device)
-    
+
     with torch.no_grad():
         logits = model(input_ids=input_ids, attention_mask=attention_mask)
         pred = torch.argmax(logits, dim=1).cpu().numpy()[0]
-    
+
     category = product_encoder.inverse_transform([pred])[0]
     return category
 
-# Gradio Interface
+# Gradio interface
 iface = gr.Interface(
     fn=predict_category,
     inputs=[
@@ -82,9 +82,9 @@ iface = gr.Interface(
     ],
     outputs="text",
     title="Product Category Predictor",
-    description="Enter the date and hazard text details to get the predicted product category."
+    description="Enter date and text details to predict the product category.",
 )
 
-# Launch app
+# Run the app
 if __name__ == "__main__":
     iface.launch()