# streamlit_app.py

import streamlit as st
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from pathlib import Path
import pandas as pd

# Local model path
MODEL_DIR = Path(__file__).parent

# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model.to(device)
model.eval()

# Inference function
def classify(prompt: str):
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = F.softmax(logits, dim=-1).squeeze().cpu()
        pred = torch.argmax(probs).item()
        confidence = probs[pred].item()
        return pred, confidence

# Streamlit UI
st.title("QDG Classifier by DEJAN")
user_input = st.text_area("Enter one prompt per line:")

if st.button("Classify"):
    prompts = [line.strip() for line in user_input.strip().split("\n") if line.strip()]
    if not prompts:
        st.warning("Please enter at least one prompt.")
    else:
        records = []
        with st.spinner("Classifying..."):
            for p in prompts:
                label, conf = classify(p)
                records.append({"Prompt": p, "Grounding": label, "Confidence": round(conf, 4)})

        df = pd.DataFrame(records)
        st.dataframe(df, use_container_width=True)