Update src/predict.py

src/predict.py

@@ -328,37 +328,15 @@ def test_predictions():
         actual = "DRUG" if label == 1 else "NON_DRUG"
         logger.info(f"Expected: {expected}, Got: {actual}, Probability: {prob:.4f}")
 
+# ===========================
+# HF Spaces / Production Ready
+# ===========================
+
 if __name__ == "__main__":
-    # Load model once when script is run directly
+    # Load model once
     load_model()
-    
-    # Run tests
+
+    # Optional: run test predictions to verify setup
     test_predictions()
     
-    # Interactive testing
-    print("\n" + "="*50)
-    print("Interactive Drug Detection Testing")
-    print("Enter text to classify (or 'quit' to exit)")
-    print("="*50)
-    
-    while True:
-        try:
-            user_input = input("\nEnter text: ").strip()
-            if user_input.lower() in ['quit', 'exit', 'q']:
-                break
-            
-            if user_input:
-                label, prob = predict(user_input)
-                result = "🚨 DRUG" if label == 1 else "✅ NON_DRUG"
-                confidence = max(prob, 1-prob)
-                print(f"Result: {result}")
-                print(f"Drug Probability: {prob*100:.2f}%")
-                print(f"Confidence: {confidence*100:.2f}%")
-            else:
-                print("Please enter some text.")
-                
-        except KeyboardInterrupt:
-            print("\nExiting...")
-            break
-        except Exception as e:
-            print(f"Error: {e}")
+    # No interactive input() calls here – compatible with Hugging Face Spaces