import streamlit as st
import torch
import joblib  # If you saved the scaler
import dill

# Load the model
loaded_model = torch.load("iris_ann_full_model.pth", pickle_module=dill)
loaded_model.eval()  # Set to evaluation mode

# Load the scaler
scaler = joblib.load("scaler.pkl")  # Ensure scaler is loaded

# Create a new sample and transform it
sample = torch.tensor([scaler.transform([[5.1, 3.5, 1.4, 0.2]])[0]], dtype=torch.float32)

# Ensure model and input are on the same device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
loaded_model.to(device)
sample = sample.to(device)

# Predict the class
with torch.no_grad():
    output = loaded_model(sample)  # Get raw output
    _, predicted_class = torch.max(output, 1)  # Get class index

# Display result in Streamlit
st.write(f"Predicted class: {predicted_class.item()}")