import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

# Load model & tokenizer
model_path = "./biobert_model"  # or "your-username/your-model-name" if from Hugging Face Hub
model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.eval()

labels = ["No interaction", "Mild", "Moderate", "Severe"]

def predict_interaction(text):
    encoding = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**encoding)
        probs = F.softmax(outputs.logits, dim=1)
        pred = torch.argmax(probs, dim=1).item()
    return f"🧠 Prediction: {labels[pred]}"

gr.Interface(fn=predict_interaction, inputs="text", outputs="text", title="Drug Interaction Predictor").launch()