Enhance prediction function with validation checks and improved error handling

utils/prediction.py

@@ -46,43 +46,43 @@ def predict_sentence(model, sentence, tokenizer, label_encoder):
     """
     Make prediction for a single sentence with label validation.
     """
+    # Validation checks
+    if model is None:
+        print("Error: Model not loaded")
+        return "Error: Model not loaded", 0.0
+    if tokenizer is None:
+        print("Error: Tokenizer not loaded")
+        return "Error: Tokenizer not loaded", 0.0
+    if label_encoder is None:
+        print("Error: Label encoder not loaded")
+        return "Error: Label encoder not loaded", 0.0
+        
+    # Force CPU device
+    device = torch.device('cpu')
+    model = model.to(device)
     model.eval()
     
     # Tokenize
-    encoding = tokenizer(
-        sentence,
-        add_special_tokens=True,
-        max_length=512,
-        padding='max_length',
-        truncation=True,
-        return_tensors='pt'
-    )
-    
     try:
+        encoding = tokenizer(
+            sentence,
+            add_special_tokens=True,
+            max_length=512,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt'
+        ).to(device)
+        
         with torch.no_grad():
-            # Get model outputs
             outputs = model(encoding['input_ids'], encoding['attention_mask'])
             probabilities = torch.softmax(outputs, dim=1)
-            
-            # Get prediction and probability
             prob, pred_idx = torch.max(probabilities, dim=1)
+            predicted_label = label_encoder.classes_[pred_idx.item()]
+            return predicted_label, prob.item()
             
-            # Validate prediction index
-            if pred_idx.item() >= len(label_encoder.classes_):
-                print(f"Warning: Model predicted invalid label index {pred_idx.item()}")
-                return "Unknown", 0.0
-                
-            # Convert to label
-            try:
-                predicted_class = label_encoder.classes_[pred_idx.item()]
-                return predicted_class, prob.item()
-            except IndexError:
-                print(f"Warning: Invalid label index {pred_idx.item()}")
-                return "Unknown", 0.0
-                
     except Exception as e:
         print(f"Prediction error: {str(e)}")
-        return "Error", 0.0
+        return f"Error: {str(e)}", 0.0
     
 def print_labels(label_encoder, show_counts=False):
     """Print all labels and their corresponding indices"""