import streamlit as st
from transformers import pipeline

# Load the Hugging Face model for drug interaction prediction
@st.cache_resource
def load_model():
    # Use a different model that is available
    model = pipeline("text-classification", model="dmis-lab/biobert-base-cased-v1.1")
    return model

# Main function to interact with the Streamlit app
def main():
    st.title("💊 Drug Interaction Predictor")
    st.write("Enter the names of drugs to predict potential interactions.")

    # Input fields for drug names
    drug1 = st.text_input("Enter Drug 1:")
    drug2 = st.text_input("Enter Drug 2:")
    drug3 = st.text_input("Enter Drug 3 (optional):")

    # Load the model
    model = load_model()

    # Check interactions when the button is clicked
    if st.button("Check Interactions"):
        if drug1 or drug2 or drug3:
            drugs = [drug for drug in [drug1, drug2, drug3] if drug]
            st.write("### Checking interactions...")

            interactions = []

            # Predict interaction for each pair of drugs
            for i in range(len(drugs)):
                for j in range(i+1, len(drugs)):
                    input_text = f"{drugs[i]} interacts with {drugs[j]}"
                    prediction = model(input_text)
                    label = prediction[0]['label']

                    if label == "1":
                        interactions.append(f"⚠️ {drugs[i]} and {drugs[j]} have a potential interaction.")
                    else:
                        interactions.append(f"✅ No significant interaction between {drugs[i]} and {drugs[j]}.")

            if interactions:
                for interaction in interactions:
                    st.write(interaction)
        else:
            st.warning("Please enter at least one drug.")
            
if __name__ == "__main__":
    main()