Update src/predict.py

src/predict.py

@@ -268,7 +268,7 @@ def validate_and_fix_config(model_path: str) -> bool:
 # Enhanced model loading with multiple fallback strategies
 # =======================
 def load_model_with_fallback(model_name: str) -> bool:
-    """Enhanced model loading with multiple fallback strategies"""
+    """Simplified model loading for HF Spaces"""
     global model, tokenizer, model_loaded
     
     with model_lock:
@@ -277,110 +277,40 @@ def load_model_with_fallback(model_name: str) -> bool:
         
         logger.info(f"Loading model: {model_name}")
         
-        # Strategy 1: Load local model with Auto classes (most compatible)
-        if os.path.exists(model_name):
-            try:
-                logger.info("Strategy 1: Loading with Auto classes...")
+        try:
+            # Check if local model exists
+            if os.path.exists(model_name):
+                logger.info("Loading local model...")
                 tokenizer = AutoTokenizer.from_pretrained(
                     model_name, 
                     use_fast=True,
-                    do_lower_case=True
+                    do_lower_case=True,
+                    local_files_only=True  # Force local loading
                 )
                 model = AutoModelForSequenceClassification.from_pretrained(
                     model_name,
                     num_labels=2,
-                    ignore_mismatched_sizes=True
-                )
-                model.to(device)
-                model.eval()
-                logger.info("✅ Successfully loaded with Auto classes")
-                model_loaded = True
-                return True
-            except Exception as e:
-                logger.error(f"Strategy 1 failed: {e}")
-        
-        # Strategy 2: Load with DistilBERT classes
-        if os.path.exists(model_name):
-            try:
-                logger.info("Strategy 2: Loading with DistilBERT classes...")
-                validate_and_fix_config(model_name)
-                
-                tokenizer = DistilBertTokenizerFast.from_pretrained(
-                    model_name,
-                    do_lower_case=True
-                )
-                model = DistilBertForSequenceClassification.from_pretrained(
-                    model_name,
-                    ignore_mismatched_sizes=True
+                    ignore_mismatched_sizes=True,
+                    local_files_only=True  # Force local loading
                 )
-                model.to(device)
-                model.eval()
-                logger.info("✅ Successfully loaded with DistilBERT classes")
-                model_loaded = True
-                return True
-            except Exception as e:
-                logger.error(f"Strategy 2 failed: {e}")
-        
-        # Strategy 3: Create model with custom config + load weights
-        if os.path.exists(model_name):
-            try:
-                logger.info("Strategy 3: Loading with custom configuration...")
-                config = DistilBertConfig(
-                    vocab_size=30522,
-                    max_position_embeddings=512,
-                    dim=768,
-                    n_layers=6,
-                    n_heads=12,
-                    hidden_dim=3072,
-                    dropout=0.1,
-                    attention_dropout=0.1,
-                    activation='gelu',
-                    num_labels=2,
-                    id2label={0: "NON_DRUG", 1: "DRUG"},
-                    label2id={"NON_DRUG": 0, "DRUG": 1}
-                )
-                
-                tokenizer = DistilBertTokenizerFast.from_pretrained(
+            else:
+                # Fallback to a working pre-trained model
+                logger.warning("Local model not found, using fallback...")
+                tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
+                model = AutoModelForSequenceClassification.from_pretrained(
                     'distilbert-base-uncased',
-                    do_lower_case=True
+                    num_labels=2
                 )
-                model = DistilBertForSequenceClassification(config)
-                
-                # Try to load weights
-                weights_path = os.path.join(model_name, "pytorch_model.bin")
-                if os.path.exists(weights_path):
-                    state_dict = torch.load(weights_path, map_location=device)
-                    model.load_state_dict(state_dict, strict=False)
-                    logger.info("✅ Loaded custom weights")
-                
-                model.to(device)
-                model.eval()
-                logger.info("✅ Successfully loaded with custom config")
-                model_loaded = True
-                return True
-            except Exception as e:
-                logger.error(f"Strategy 3 failed: {e}")
-        
-        # Strategy 4: Use pre-trained DistilBERT as fallback
-        try:
-            logger.warning("Strategy 4: Falling back to pre-trained DistilBERT...")
-            tokenizer = DistilBertTokenizerFast.from_pretrained(
-                'distilbert-base-uncased',
-                do_lower_case=True
-            )
-            model = DistilBertForSequenceClassification.from_pretrained(
-                'distilbert-base-uncased',
-                num_labels=2
-            )
+            
             model.to(device)
             model.eval()
-            logger.warning("⚠️ Using pre-trained DistilBERT (not fine-tuned for drug detection)")
             model_loaded = True
+            logger.info("Model loaded successfully")
             return True
+            
         except Exception as e:
-            logger.error(f"All strategies failed: {e}")
+            logger.error(f"Model loading failed: {e}")
             return False
-
 # =======================
 # Optimized prediction function with enhanced accuracy
 # =======================