6 commit

app.py

@@ -47,6 +47,17 @@ class MultilabelBertClassifier(nn.Module):
         pooled_output = outputs.last_hidden_state[:, 0, :]  # Use [CLS] token
         return self.classifier(pooled_output)
 
+# Function to remove 'module.' prefix from state dict keys
+def remove_module_prefix(state_dict):
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        if key.startswith('module.'):
+            new_key = key[7:]  # Remove 'module.' prefix
+        else:
+            new_key = key
+        new_state_dict[new_key] = value
+    return new_state_dict
+
 # Load model function
 @st.cache_resource
 def load_model():
@@ -60,16 +71,28 @@ def load_model():
     model_path = "model/alergen_model.pt"
     
     if os.path.exists(model_path):
-        # Load model weights
-        checkpoint = torch.load(model_path, map_location=device)
-        if 'model_state_dict' in checkpoint:
-            model.load_state_dict(checkpoint['model_state_dict'])
-        else:
-            model.load_state_dict(checkpoint)
-        
-        model.to(device)
-        model.eval()
-        return model, tokenizer
+        try:
+            # Load model weights
+            checkpoint = torch.load(model_path, map_location=device)
+            
+            # Check if state_dict is directly in checkpoint or under 'model_state_dict' key
+            if 'model_state_dict' in checkpoint:
+                state_dict = checkpoint['model_state_dict']
+            else:
+                state_dict = checkpoint
+                
+            # Remove 'module.' prefix if it exists
+            state_dict = remove_module_prefix(state_dict)
+            
+            # Load the processed state dict
+            model.load_state_dict(state_dict)
+            
+            model.to(device)
+            model.eval()
+            return model, tokenizer
+        except Exception as e:
+            st.error(f"Error loading model: {str(e)}")
+            return None, tokenizer
     else:
         st.error("Model file not found. Please upload the model file.")
         return None, tokenizer
@@ -111,10 +134,13 @@ def main():
     st.title("🍲 Allergen Detection in Indonesian Recipes")
     st.write("This app predicts common allergens in your recipe based on ingredients.")
     
+    # Create directory for model if it doesn't exist
+    os.makedirs("model", exist_ok=True)
+    
     # Sidebar for model upload
     with st.sidebar:
         st.header("Model Settings")
-        uploaded_model = st.file_uploader("Upload model file (model/alergen_model.pt)", type=["pt"])
+        uploaded_model = st.file_uploader("Upload model file (alergen_model.pt)", type=["pt"])
         
         if uploaded_model:
             # Save uploaded model
@@ -212,7 +238,7 @@ def main():
     # Instructions and information
     with st.expander("How to Use"):
         st.write("""
-        1. First, upload the trained model file (`model/alergen_model.pt`) using the sidebar uploader
+        1. First, upload the trained model file (`alergen_model.pt`) using the sidebar uploader
         2. Enter your recipe ingredients in the text box (in Indonesian)
         3. Click the "Detect Allergens" button to analyze the recipe
         4. View the results showing which allergens are present in your recipe