Refactor model loading function to handle checkpoints and improve error handling

utils/prediction.py

@@ -19,57 +19,40 @@ def load_model_for_prediction():
             dropout=0.5
         )
         
+        # Initialize model
         model = BiLSTMAttentionBERT(config)
+        
+        # Load checkpoint
         model_path = hf_hub_download(
             repo_id="joko333/BiLSTM_v01",
             filename="model_epoch8_acc72.53.pt"
         )
-        state_dict = torch.load(model_path, map_location='cpu')
-        model.load_state_dict(state_dict)
+        checkpoint = torch.load(model_path, map_location='cpu')
         
-        # Test Hugging Face connectivity
-        st.write("Testing connection to Hugging Face...")
-        response = requests.get("https://huggingface.co/joko333/BiLSTM_v01")
-        if response.status_code != 200:
-            st.error(f"Cannot connect to Hugging Face. Status code: {response.status_code}")
+        # Extract model state dict from checkpoint
+        if 'model_state_dict' in checkpoint:
+            state_dict = checkpoint['model_state_dict']
+            model.load_state_dict(state_dict)
+            st.write("Model loaded successfully")
+        else:
+            st.error("Invalid checkpoint format")
             return None, None, None
             
-        # Load model with logging
-        st.write("Loading BiLSTM model...")
-        model = BiLSTMAttentionBERT.from_pretrained(
-            "joko333/BiLSTM_v01",
-            hidden_dim=128,
-            num_classes=22,
-            num_layers=2,
-            dropout=0.5
-        )
-        st.write("Model loaded successfully")
-        
-        # Initialize label encoder
-        st.write("Initializing label encoder...")
+        # Initialize label encoder from checkpoint
         label_encoder = LabelEncoder()
-        label_encoder.classes_ = np.array(['Addition', 'Causal', 'Cause and Effect', 
-                                         'Clarification', 'Comparison', 'Concession', 
-                                         'Conditional', 'Contrast', 'Contrastive Emphasis',
-                                         'Definition', 'Elaboration', 'Emphasis', 
-                                         'Enumeration', 'Explanation', 'Generalization', 
-                                         'Illustration', 'Inference', 'Problem Solution', 
-                                         'Purpose', 'Sequential', 'Summary', 
-                                         'Temporal Sequence'])
-        st.write("Label encoder initialized")
-        
+        if 'label_encoder_classes' in checkpoint:
+            label_encoder.classes_ = checkpoint['label_encoder_classes']
+        else:
+            st.error("Label encoder data not found in checkpoint")
+            return None, None, None
+            
         # Load tokenizer
-        st.write("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained('dmis-lab/biobert-base-cased-v1.2')
-        st.write("Tokenizer loaded successfully")
         
         return model, label_encoder, tokenizer
         
     except Exception as e:
-        st.error(f"Detailed error: {str(e)}")
-        st.error(f"Error type: {type(e).__name__}")
-        import traceback
-        st.error(f"Traceback: {traceback.format_exc()}")
+        st.error(f"Error loading model: {str(e)}")
         return None, None, None
 
 def predict_sentence(model, sentence, tokenizer, label_encoder):