File size: 3,896 Bytes
42e236f eeb6ce3 42e236f eda24bc 42e236f eeb6ce3 42e236f eeb6ce3 42e236f eeb6ce3 42e236f eeb6ce3 42e236f |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
# app.py
import streamlit as st
import requests
import os
from dotenv import load_dotenv
dotenv_path = os.path.join(os.path.dirname(__file__), "..", ".env")
load_dotenv(dotenv_path=dotenv_path)
API_TOKEN = os.getenv("HF_TOKEN")
MODEL_ID = "facebook/bart-large-mnli"
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
headers = {"Authorization": f"Bearer {API_TOKEN}"}
def query_hf_api(payload):
"""
Sends a request to the Hugging Face Inference API.
"""
response = requests.post(API_URL, headers=headers, json=payload)
response.raise_for_status()
return response.json()
st.set_page_config(layout="wide", page_title="Zero-Shot Text Classifier")
st.title("🎯 Zero-Shot Text Classifier")
st.markdown(
f"""
This app uses the `{MODEL_ID}` Zero-Shot Classification model from Hugging Face
to classify text into categories you define, even if the model hasn't been explicitly trained on those categories.
Enter a piece of text and a comma-separated list of potential labels.
"""
)
text_to_classify = st.text_area(
"Enter text to classify:",
"The new AI regulations will have a significant impact on technology startups.",
height=100,
)
candidate_labels_input = st.text_input(
"Enter candidate labels (comma-separated):",
"politics, business, technology, sports, education, finance",
)
if st.button("Classify Text"):
if not API_TOKEN:
st.error(
"Hugging Face API token not found. Please set HUGGING_FACE_API_TOKEN in your .env file."
)
elif not text_to_classify.strip():
st.warning("Please enter some text to classify.")
elif not candidate_labels_input.strip():
st.warning("Please enter some candidate labels.")
else:
candidate_labels = [
label.strip()
for label in candidate_labels_input.split(",")
if label.strip()
]
if not candidate_labels:
st.warning(
"No valid candidate labels provided after parsing. Please ensure labels are separated by commas and are not empty."
)
else:
payload = {
"inputs": text_to_classify,
"parameters": {"candidate_labels": candidate_labels},
}
with st.spinner(f"Asking the AI model ({MODEL_ID})..."):
try:
result = query_hf_api(payload)
st.subheader("Classification Results:")
if (
isinstance(result, dict)
and "labels" in result
and "scores" in result
):
results_with_scores = list(
zip(result["labels"], result["scores"])
)
for label, score in results_with_scores:
st.write(f"**Label:** {label} - **Score:** {score:.4f}")
elif isinstance(result, dict) and "error" in result:
st.error(f"API Error: {result['error']}")
if "estimated_time" in result:
st.info(
f"The model might be loading. Estimated time: {result['estimated_time']:.2f} seconds. Please try again shortly."
)
else:
st.error("Received an unexpected response format from the API.")
st.json(result)
except requests.exceptions.RequestException as e:
st.error(f"API Request Failed: {e}")
if e.response is not None:
st.error(f"Response content: {e.response.text}")
except Exception as e:
st.error(f"An unexpected error occurred: {e}")
st.markdown("---")
|