Update app.py

app.py

@@ -1,16 +1,19 @@
+import os
 import pickle
 import requests
+from huggingface_hub import hf_hub_download
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+import torch.nn.functional as F
 import gradio as gr
 
-# Load label encoder
-label_encoder_path = 'Fredaaaaaa/hybrid_model/label_encoder.pkl'  # Path to the uploaded label encoder
+# Function to download the label encoder from Hugging Face
+label_encoder_path = hf_hub_download(repo_id="Fredaaaaaa/hybrid_model", filename="label_encoder.pkl")
 with open(label_encoder_path, 'rb') as f:
     label_encoder = pickle.load(f)
 
-# Load model and tokenizer from Hugging Face model directory
-model_name = "Fredaaaaaa/hybrid_model"
+# Load the model and tokenizer
+model_name = "Fredaaaaaa/hybrid_model"  # Update with the correct model repo
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -18,7 +21,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 def get_drug_features(drug1, drug2):
     api_url = f'https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{drug1}/property/SMILES,Pharmacology,Toxicity,Mechanism-of-action,Route-of-elimination,Metabolism/JSON'
     response = requests.get(api_url)
-    
+
     if response.status_code == 200:
         data = response.json()
         drug_features = {
@@ -31,48 +34,72 @@ def get_drug_features(drug1, drug2):
         }
         return drug_features
     else:
-        return None
+        return None  # If no data is returned, handle the missing values gracefully
 
-# Function to predict severity based on drug interactions
-def predict_severity(drug1, drug2):
-    # Check if drugs are present in the dataset
+# Function to check if the drugs are in the dataset
+def check_drugs_in_dataset(drug1, drug2, df):
     if drug1 in df['Drug 1_normalized'].values and drug2 in df['Drug 2_normalized'].values:
-        # Predict severity based on the dataset (your logic for prediction)
-        # Add your prediction logic based on the dataset here
-        severity = "Predicted severity based on dataset logic"
-        return severity
+        return True
+    return False
+
+# Function to process the input and predict severity
+def predict_severity(drug1, drug2, df):
+    # Check if drugs are in the dataset first
+    if check_drugs_in_dataset(drug1, drug2, df):
+        # Preprocess text data (interaction description)
+        interaction_description = f"{drug1} interacts with {drug2}"  # Example interaction description
+        inputs = tokenizer(interaction_description, return_tensors="pt", padding=True, truncation=True, max_length=128)
+        input_ids = inputs['input_ids'].to(device)
+        attention_mask = inputs['attention_mask'].to(device)
+
+        # Run the model to get predictions
+        model.eval()
+        with torch.no_grad():
+            outputs = model(input_ids, attention_mask=attention_mask)
+        
+        # Get the predicted class
+        prediction = torch.argmax(outputs.logits, dim=1).item()
+
+        # Map the predicted class index to the severity label
+        severity_labels = ["No interaction", "Mild", "Moderate", "Severe"]
+        return severity_labels[prediction]
+    
     else:
-        # If drugs not found in dataset, call the external API to fetch their features
+        # Fetch drug features from the external API if not in the dataset
         drug_features = get_drug_features(drug1, drug2)
 
         if not drug_features:
             return "Error: Could not fetch drug features from the API."
 
         # Preprocess text data (interaction description)
-        interaction_description = f"{drug1} interacts with {drug2}"
+        interaction_description = f"{drug1} interacts with {drug2}"  # Example interaction description
         inputs = tokenizer(interaction_description, return_tensors="pt", padding=True, truncation=True, max_length=128)
-        input_ids = inputs['input_ids'].to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
-        attention_mask = inputs['attention_mask'].to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        input_ids = inputs['input_ids'].to(device)
+        attention_mask = inputs['attention_mask'].to(device)
 
         # Process drug features (SMILES, pharmacology, toxicity, etc.)
         drug_feature_values = [drug_features['SMILES'], drug_features['pharmacology'], drug_features['toxicity']]
-        drug_features_tensor = torch.tensor(drug_feature_values, dtype=torch.float32).unsqueeze(0).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        drug_features_tensor = torch.tensor(drug_feature_values, dtype=torch.float32).unsqueeze(0).to(device)
 
         # Run the model to get predictions
         model.eval()
         with torch.no_grad():
-            outputs = model(input_ids, attention_mask)
+            outputs = model(input_ids, attention_mask, drug_features_tensor)
         
         # Get the predicted class
-        prediction = torch.argmax(outputs.logits, dim=1).item()
+        prediction = torch.argmax(outputs, dim=1).item()
 
         # Map the predicted class index to the severity label
         severity_labels = ["No interaction", "Mild", "Moderate", "Severe"]
         return severity_labels[prediction]
 
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"🚀 Device: {device}")
+
 # Gradio Interface
 interface = gr.Interface(
-    fn=predict_severity,
+    fn=lambda drug1, drug2: predict_severity(drug1, drug2, df),  # Pass dataset to the prediction function
     inputs=[gr.Textbox(label="Drug 1"), gr.Textbox(label="Drug 2")],
     outputs="text",
     live=True
@@ -80,3 +107,4 @@ interface = gr.Interface(
 
 # Launch the interface
 interface.launch()
+