Update inference.py

inference.py

@@ -1,8 +1,21 @@
-import torch
-from transformers import AutoTokenizer, AutoModel
-import joblib
-from huggingface_hub import hf_hub_download
-import json
+# First try to import with fallbacks
+try:
+    import torch
+    from transformers import AutoTokenizer, AutoModel
+    import joblib
+    from huggingface_hub import hf_hub_download
+    import json
+except ImportError as e:
+    print(f"Import error: {e}")
+    # Try to install missing packages (this might not work in Spaces)
+    import subprocess
+    import sys
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "torch", "transformers", "joblib", "huggingface-hub"])
+    import torch
+    from transformers import AutoTokenizer, AutoModel
+    import joblib
+    from huggingface_hub import hf_hub_download
+    import json
 
 class DrugInteractionClassifier(torch.nn.Module):
     def __init__(self, n_classes, bert_model_name="emilyalsentzer/Bio_ClinicalBERT"):
@@ -23,38 +36,53 @@ class DrugInteractionClassifier(torch.nn.Module):
 class DDIPredictor:
     def __init__(self, repo_id="Fredaaaaaa/drug_interaction_severity"):
         self.repo_id = repo_id
+        print(f"🚀 Loading model from: {repo_id}")
         
-        # Download model files from Hugging Face
-        self.config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
-        self.model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
-        self.label_encoder_path = hf_hub_download(repo_id=repo_id, filename="label_encoder.joblib")
-        
-        # Load config
-        with open(self.config_path, "r") as f:
-            self.config = json.load(f)
-        
-        # Load tokenizer from repo
-        self.tokenizer = AutoTokenizer.from_pretrained(repo_id)
-        
-        # Load label encoder
-        self.label_encoder = joblib.load(self.label_encoder_path)
-        
-        # Initialize model
-        self.model = DrugInteractionClassifier(
-            n_classes=self.config["num_labels"],
-            bert_model_name=self.config["bert_model_name"]
-        )
-        
-        # Load weights
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.load_state_dict(
-            torch.load(self.model_path, map_location=device, weights_only=True)
-        )
-        self.model.to(device)
-        self.model.eval()
-        
-        self.device = device
-        print(f"✅ Model loaded successfully from {repo_id} on {device}")
+        try:
+            # Download model files from Hugging Face
+            print("📥 Downloading config.json...")
+            self.config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
+            
+            print("📥 Downloading pytorch_model.bin...")
+            self.model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
+            
+            print("📥 Downloading label_encoder.joblib...")
+            self.label_encoder_path = hf_hub_download(repo_id=repo_id, filename="label_encoder.joblib")
+            
+            # Load config
+            with open(self.config_path, "r") as f:
+                self.config = json.load(f)
+            
+            # Load tokenizer from repo
+            print("🔤 Loading tokenizer...")
+            self.tokenizer = AutoTokenizer.from_pretrained(repo_id)
+            
+            # Load label encoder
+            print("🏷️ Loading label encoder...")
+            self.label_encoder = joblib.load(self.label_encoder_path)
+            
+            # Initialize model
+            print("🧠 Initializing model...")
+            self.model = DrugInteractionClassifier(
+                n_classes=self.config["num_labels"],
+                bert_model_name=self.config["bert_model_name"]
+            )
+            
+            # Load weights
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            print(f"⚙️ Loading weights on {device}...")
+            self.model.load_state_dict(
+                torch.load(self.model_path, map_location=device)
+            )
+            self.model.to(device)
+            self.model.eval()
+            
+            self.device = device
+            print(f"✅ Model loaded successfully from {repo_id} on {device}")
+            
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            raise e
     
     def predict(self, text, confidence_threshold=0.0):
         """Predict drug interaction severity"""
@@ -69,7 +97,7 @@ class DDIPredictor:
             # Tokenize
             inputs = self.tokenizer(
                 text,
-                max_length=self.config["max_length"],
+                max_length=self.config.get("max_length", 128),
                 padding=True,
                 truncation=True,
                 return_tensors="pt"
@@ -104,14 +132,11 @@ class DDIPredictor:
                 "probabilities": {label: 0.0 for label in self.label_encoder.classes_}
             }
 
-# Simple test
-if __name__ == "__main__":
-    try:
-        predictor = DDIPredictor("Fredaaaaaa/drug_interaction_severity")
-        test_text = "Drug interaction may increase bleeding risk"
-        result = predictor.predict(test_text)
-        print("✅ Test successful!")
-        print(f"Prediction: {result['prediction']}")
-        print(f"Confidence: {result['confidence']:.3f}")
-    except Exception as e:
-        print(f"❌ Error: {e}")
+# Global predictor instance
+try:
+    predictor = DDIPredictor("Fredaaaaaa/drug_interaction_severity")
+    MODEL_LOADED = True
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    predictor = None
+    MODEL_LOADED = False