Update app.py

app.py

@@ -1,15 +1,137 @@
 import gradio as gr
 import requests
 import time
+import json
+import joblib
+from huggingface_hub import hf_hub_download
+import numpy as np
 
-# Try to import the predictor
+# Try to import torch with fallback
 try:
-    from inference import predictor, MODEL_LOADED
-    print("✅ Inference module imported successfully")
-except ImportError as e:
-    print(f"❌ Failed to import inference: {e}")
+    import torch
+    from transformers import AutoTokenizer, AutoModel
+    TORCH_AVAILABLE = True
+except ImportError:
+    print("Torch not available, using mock mode")
+    TORCH_AVAILABLE = False
+
+class MockPredictor:
+    """Mock predictor for when torch is not available"""
+    def __init__(self):
+        self.classes = ["Mild", "Moderate", "No Interaction", "Severe"]
+    
+    def predict(self, text):
+        return {
+            "prediction": "Moderate",
+            "confidence": 0.75,
+            "probabilities": {cls: 0.25 for cls in self.classes}
+        }
+
+# Initialize predictor
+if TORCH_AVAILABLE:
+    try:
+        # Define model class
+        class DrugInteractionClassifier(torch.nn.Module):
+            def __init__(self, n_classes, bert_model_name="emilyalsentzer/Bio_ClinicalBERT"):
+                super(DrugInteractionClassifier, self).__init__()
+                self.bert = AutoModel.from_pretrained(bert_model_name)
+                self.classifier = torch.nn.Sequential(
+                    torch.nn.Linear(self.bert.config.hidden_size, 256),
+                    torch.nn.ReLU(),
+                    torch.nn.Dropout(0.3),
+                    torch.nn.Linear(256, n_classes)
+                )
+
+            def forward(self, input_ids, attention_mask):
+                bert_output = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+                pooled_output = bert_output[0][:, 0, :]
+                return self.classifier(pooled_output)
+
+        # Load model components
+        print("Downloading model files from Fredaaaaaa/drug_interaction_severity...")
+        
+        # Download config
+        config_path = hf_hub_download(repo_id="Fredaaaaaa/drug_interaction_severity", filename="config.json")
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        
+        # Download label encoder
+        label_encoder_path = hf_hub_download(repo_id="Fredaaaaaa/drug_interaction_severity", filename="label_encoder.joblib")
+        label_encoder = joblib.load(label_encoder_path)
+        
+        # Download model weights
+        model_path = hf_hub_download(repo_id="Fredaaaaaa/drug_interaction_severity", filename="pytorch_model.bin")
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained("Fredaaaaaa/drug_interaction_severity")
+        
+        # Initialize model
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = DrugInteractionClassifier(n_classes=len(label_encoder.classes_))
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model.to(device)
+        model.eval()
+        
+        print("✅ Model loaded successfully!")
+        
+        class RealPredictor:
+            def __init__(self, model, tokenizer, label_encoder, device, config):
+                self.model = model
+                self.tokenizer = tokenizer
+                self.label_encoder = label_encoder
+                self.device = device
+                self.config = config
+            
+            def predict(self, text):
+                try:
+                    # Tokenize
+                    inputs = self.tokenizer(
+                        text,
+                        max_length=self.config.get("max_length", 128),
+                        padding=True,
+                        truncation=True,
+                        return_tensors="pt"
+                    )
+                    
+                    inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                    
+                    # Predict
+                    with torch.no_grad():
+                        outputs = self.model(inputs["input_ids"], inputs["attention_mask"])
+                        probabilities = torch.softmax(outputs, dim=1)
+                        confidence, predicted_idx = torch.max(probabilities, dim=1)
+                    
+                    predicted_label = self.label_encoder.inverse_transform([predicted_idx.item()])[0]
+                    
+                    # Get all probabilities
+                    all_probs = {
+                        self.label_encoder.inverse_transform([i])[0]: prob.item() 
+                        for i, prob in enumerate(probabilities[0])
+                    }
+                    
+                    return {
+                        "prediction": predicted_label,
+                        "confidence": confidence.item(),
+                        "probabilities": all_probs
+                    }
+                    
+                except Exception as e:
+                    return {
+                        "prediction": f"Error: {str(e)}",
+                        "confidence": 0.0,
+                        "probabilities": {label: 0.0 for label in self.label_encoder.classes_}
+                    }
+        
+        predictor = RealPredictor(model, tokenizer, label_encoder, device, config)
+        MODEL_LOADED = True
+        
+    except Exception as e:
+        print(f"Error loading real model: {e}")
+        predictor = MockPredictor()
+        MODEL_LOADED = False
+else:
+    predictor = MockPredictor()
     MODEL_LOADED = False
-    predictor = None
 
 def fetch_pubchem_data(drug_name):
     """Fetch drug data from PubChem by name"""
@@ -26,7 +148,7 @@ def fetch_pubchem_data(drug_name):
         
         cid = search_response.json()['IdentifierList']['CID'][0]
         
-        # Fetch basic compound data
+        # Fetch compound data
         compound_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/CanonicalSMILES,MolecularWeight,IUPACName/JSON"
         compound_response = requests.get(compound_url, timeout=10)
         
@@ -51,10 +173,12 @@ def generate_interaction_description(drug1_data, drug2_data):
         if mw1 and mw2:
             mw_diff = abs(mw1 - mw2)
             if mw_diff > 300:
-                descriptions.append("Significant molecular size difference")
+                descriptions.append("Significant molecular size difference may affect metabolism")
+            elif mw_diff > 100:
+                descriptions.append("Moderate molecular size difference")
         
         if not descriptions:
-            descriptions.append("Potential drug interaction")
+            descriptions.append("Potential pharmacokinetic interaction")
         
         return ". ".join(descriptions) + ". Clinical evaluation recommended."
     except:
@@ -63,9 +187,6 @@ def generate_interaction_description(drug1_data, drug2_data):
 def predict_ddi(drug1_name, drug2_name):
     """Main prediction function"""
     try:
-        if not MODEL_LOADED or predictor is None:
-            return "Model not loaded. Please check requirements.txt", "", "", ""
-        
         if not drug1_name or not drug2_name:
             return "Please enter both drug names", "", "", ""
         
@@ -84,36 +205,67 @@ def predict_ddi(drug1_name, drug2_name):
         
         # Prepare output
         drug_info = f"""
-        **{drug1_name}**: MW={drug1_data.get('MolecularWeight', 'N/A')} g/mol
-        **{drug2_name}**: MW={drug2_data.get('MolecularWeight', 'N/A')} g/mol
+        **{drug1_name}**: 
+        - Molecular Weight: {drug1_data.get('MolecularWeight', 'N/A')} g/mol
+        - IUPAC Name: {drug1_data.get('IUPACName', 'N/A')}
+        
+        **{drug2_name}**: 
+        - Molecular Weight: {drug2_data.get('MolecularWeight', 'N/A')} g/mol
+        - IUPAC Name: {drug2_data.get('IUPACName', 'N/A')}
         """
         
         prediction_output = f"""
-        **Prediction:** {result['prediction']}
+        ## 🔍 Prediction Results
+        
+        **Severity:** **{result['prediction']}**
         **Confidence:** {result['confidence']:.1%}
-        **Description:** {interaction_description}
+        
+        **Generated Description:**
+        {interaction_description}
+        
+        **Probabilities:**
+        {', '.join([f'{k}: {v:.1%}' for k, v in result['probabilities'].items()])}
         """
         
-        return prediction_output, drug_info, interaction_description, "✅ Success"
+        status = "✅ Success" if MODEL_LOADED else "⚠️ Using mock data (torch not available)"
+        
+        return prediction_output, drug_info, interaction_description, status
         
     except Exception as e:
         return f"Error: {str(e)}", "", "", ""
 
-# Create simple interface
-with gr.Blocks(title="Drug Interaction Predictor") as demo:
-    gr.Markdown("# 💊 Drug Interaction Predictor")
-    gr.Markdown("Model: Fredaaaaaa/drug_interaction_severity")
+# Create interface
+with gr.Blocks(title="Drug Interaction Predictor", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 💊 Drug Interaction Severity Predictor")
+    gr.Markdown("Model: [Fredaaaaaa/drug_interaction_severity](https://huggingface.co/Fredaaaaaa/drug_interaction_severity)")
     
     with gr.Row():
-        drug1 = gr.Textbox(label="Drug 1", placeholder="e.g., Aspirin")
-        drug2 = gr.Textbox(label="Drug 2", placeholder="e.g., Warfarin")
+        drug1 = gr.Textbox(label="First Drug", placeholder="e.g., Warfarin", value="Warfarin")
+        drug2 = gr.Textbox(label="Second Drug", placeholder="e.g., Aspirin", value="Aspirin")
     
-    predict_btn = gr.Button("Predict", variant="primary")
+    predict_btn = gr.Button("🔬 Predict Interaction", variant="primary")
     
-    output = gr.Markdown("## Results will appear here")
-    drug_info = gr.Markdown()
-    interaction_desc = gr.Textbox(label="Generated Description")
-    status = gr.Textbox(label="Status")
+    with gr.Row():
+        output = gr.Markdown("## 📊 Results will appear here")
+    
+    with gr.Row():
+        drug_info = gr.Markdown("### 💊 Drug Properties")
+    
+    with gr.Row():
+        interaction_desc = gr.Textbox(label="📝 Generated Description", interactive=False)
+        status = gr.Textbox(label="🔄 Status", interactive=False)
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            ["Warfarin", "Aspirin"],
+            ["Simvastatin", "Clarithromycin"],
+            ["Digoxin", "Quinine"],
+            ["Metformin", "Ibuprofen"]
+        ],
+        inputs=[drug1, drug2],
+        label="💡 Example Drug Pairs"
+    )
     
     predict_btn.click(
         predict_ddi,