Refactor prediction functions to remove unnecessary device handling and clean up code

utils/prediction.py

@@ -7,9 +7,6 @@ import numpy as np
 
 
 def load_model_for_prediction():
-    # Force CPU
-    device = torch.device('cpu')
-    torch.backends.mps.enabled = False
     
     try:
         # Load model from Hugging Face Hub
@@ -19,7 +16,7 @@ def load_model_for_prediction():
             num_classes=22,
             num_layers=2,
             dropout=0.5
-        ).to(device)
+        )
         
         model.eval()
         
@@ -45,12 +42,10 @@ def load_model_for_prediction():
         print(f"Error loading model components: {str(e)}")
         return None, None, None
 
-def predict_sentence(model, sentence, tokenizer, label_encoder, device=None):
+def predict_sentence(model, sentence, tokenizer, label_encoder):
     """
     Make prediction for a single sentence with label validation.
     """
-    device = torch.device('cpu')
-    model = model.to(device)
     model.eval()
     
     # Tokenize
@@ -61,7 +56,7 @@ def predict_sentence(model, sentence, tokenizer, label_encoder, device=None):
         padding='max_length',
         truncation=True,
         return_tensors='pt'
-    ).to(device)
+    )
     
     try:
         with torch.no_grad():
@@ -98,25 +93,4 @@ def print_labels(label_encoder, show_counts=False):
     print("-" * 40)
     print(f"Total number of classes: {len(label_encoder.classes_)}\n")
 
-def predict_sentence2(sentence, model, tokenizer, label_encoder):
-    # Tokenize the input
-    inputs = tokenizer(sentence, 
-                      padding=True,
-                      truncation=True,
-                      return_tensors='pt',
-                      max_length=512)
-    
-    # Move inputs to the same device as model
-    device = next(model.parameters()).device
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    
-    # Make prediction
-    with torch.no_grad():
-        outputs = model(**inputs)
-        predictions = torch.argmax(outputs.logits, dim=1)
-    
-    # Convert prediction to label
-    predicted_label = label_encoder.inverse_transform(predictions.cpu().numpy())[0]
-    
-    return predicted_label