File size: 1,706 Bytes
a24e9d7 |
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 50 51 |
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
# Load tokenizer and model from Hugging Face model hub
model_name = "dejanseo/Intent-XS"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval() # Set the model to evaluation mode
# Human-readable labels
label_map = {
1: 'Commercial',
2: 'Non-Commercial',
3: 'Branded',
4: 'Non-Branded',
5: 'Informational',
6: 'Navigational',
7: 'Transactional',
8: 'Commercial Investigation',
9: 'Local',
10: 'Entertainment'
}
# Function to perform inference
def get_predictions(text):
inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probabilities = torch.sigmoid(logits).squeeze()
predictions = (probabilities > 0.5).int()
return probabilities.numpy(), predictions.numpy()
# Streamlit user interface
st.title('Multi-label Classification with Intent-XS')
query = st.text_input("Enter your query:")
if st.button('Submit'):
if query:
probabilities, predictions = get_predictions(query)
result = {label_map[i+1]: f"Probability: {prob:.2f}" for i, prob in enumerate(probabilities) if predictions[i] == 1}
if result:
st.write("Predicted Categories:")
for label, prob in result.items():
st.write(f"{label}: {prob}")
else:
st.write("No relevant categories predicted.")
else:
st.write("Please enter a query to get predictions.")
|