|
|
|
|
| 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
|
|
|
|
|
| MODEL_DIR = Path(__file__).parent
|
|
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
| model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
|
| model.to(device)
|
| model.eval()
|
|
|
|
|
| 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
|
|
|
|
|
| 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)
|
|
|
