Create inference.py

inference.py

@@ -0,0 +1,62 @@
+# inference.py
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from torch.utils.data import DataLoader, Dataset
+import numpy as np
+import pandas as pd
+
+class DrugInteractionDataset(Dataset):
+    """Custom Dataset for drug interaction data."""
+    def __init__(self, description, tokenizer, max_length=512):
+        self.description = [description]
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+    
+    def __len__(self):
+        return 1
+    
+    def __getitem__(self, idx):
+        encoding = self.tokenizer(
+            self.description[idx],
+            padding='max_length',
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors='pt'
+        )
+        return {
+            'input_ids': encoding['input_ids'].squeeze(),
+            'attention_mask': encoding['attention_mask'].squeeze(),
+        }
+
+class DDIPredictor:
+    def __init__(self, model_repo="Fredaaaaaa/drug_interaction_severity"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_repo)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_repo)
+        self.model.eval()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+
+    def predict(self, interaction_description):
+        # Create dataset and dataloader
+        dataset = DrugInteractionDataset(interaction_description, self.tokenizer)
+        dataloader = DataLoader(dataset, batch_size=1)
+
+        # Get prediction
+        with torch.no_grad():
+            for batch in dataloader:
+                input_ids = batch['input_ids'].to(self.device)
+                attention_mask = batch['attention_mask'].to(self.device)
+                outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+                logits = outputs.logits
+                probabilities = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+                prediction = torch.argmax(logits, dim=-1).cpu().item()
+
+        # Map prediction to label (adjust based on your training labels)
+        label_map = {0: "No Interaction", 1: "Mild", 2: "Moderate", 3: "Severe"}  # Update based on your classes
+        confidence = probabilities[prediction] * 100
+
+        return {
+            "prediction": label_map.get(prediction, "Unknown"),
+            "confidence": confidence,
+            "probabilities": {k: v for k, v in enumerate(probabilities)}
+        }