File size: 1,596 Bytes
281c498 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# 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)
|