app.py

import time
import requests
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
import streamlit as st
import pandas as pd
import numpy as np


######## ------------ APP CONFIG ------------------------------------------------------------------------------
st.set_page_config(
    page_title="Smallest boy",
    page_icon="🐻",
    initial_sidebar_state="expanded"
)

######## ------------ ACCESS TOKEN ----------------------------------------------------------------------------
ACCESS_TOKEN = st.secrets["API_KEY"]
API_URL = "https://api-inference.huggingface.co/models/MoritzLaurer/multilingual-MiniLMv2-L6-mnli-xnli"
headers = {"Authorization": f"Bearer {ACCESS_TOKEN}"}


###### ------------- RETRY OPTION ------------------------------------------------------
def create_retry_session():
    session = requests.Session()
    retries = Retry(
        total=3,  # number of retries
        backoff_factor=1,  # wait 1, 2, 4 seconds between retries
        status_forcelist=[500, 502, 503, 504],  # retry on these HTTP status codes
    )
    session.mount('https://', HTTPAdapter(max_retries=retries))
    return session


###### --------------- INITIALIZE THE MODEL -----------------------------------------------
@st.cache_data(ttl=3600)
def initialize_model():
    """Initialize the model and cache the initialization"""
    try:
        session = create_retry_session()
        dummy_payload = {
            "inputs": "test",
            "parameters": {"candidate_labels": ["test1", "test2"]}
        }
        
        with st.spinner('Initializing model... This might take a few moments'):
            response = session.post(API_URL, headers=headers, json=dummy_payload)
            result = response.json()
            
            if "error" in result and "is currently loading" in result["error"]:
                wait_time = result.get("estimated_time", 20)
                time.sleep(wait_time)
                response = session.post(API_URL, headers=headers, json=dummy_payload)
            
            return True
            
    except requests.exceptions.RequestException as e:
        st.error(f"""
        Unable to initialize model. Please try:
        1. Waiting a few minutes and refreshing the page
        2. Checking your internet connection
        3. Verifying the API is not under maintenance
        
        Error details: {str(e)}
        """)
        return False


############# --------------------- DATA ------------------------------

@st.cache_data(ttl=60)
def query(payload):
    try:
        session = create_retry_session()
        
        with st.spinner('Processing... Please wait'):
            response = session.post(API_URL, headers=headers, json=payload)
            response.raise_for_status()
            return response.json()
            
    except requests.exceptions.RequestException as e:
        error_message = str(e)
        if "503" in error_message:
            st.error("""
            The service is currently unavailable. This might be due to:
            - High server load
            - Temporary maintenance
            - Service disruption
            
            Please try again in a few minutes.
            """)
        else:
            st.error(f"API request failed: {error_message}")
        return None
    
    
# Status indicator in sidebar
st.sidebar.markdown("### Service Status")
service_status = st.sidebar.empty()

# Initialize model at startup
model_ready = initialize_model()
if model_ready:
    service_status.success("Service is running")
else:
    service_status.error("Service is unavailable")


    

# Initialize model at startup
model_ready = initialize_model()

st.title("Classify your text with your own labels 🦾")

# Input text area
user_input = st.text_area(
    "Enter text to classify:",
    "Hi, I recently bought a device from your company but it is not working as advertised and I would like to get reimbursed!"
)

# Default labels
default_labels = ["refund", "legal", "faq"]

# Label selection method
label_method = st.radio(
    "Choose label input method:",
    ["Use default labels", "Enter custom labels"]
)

if label_method == "Use default labels":
    labels = default_labels
else:
    # Custom labels input
    custom_labels = st.text_area(
        "Enter custom labels (one per line):",
        "\n".join(default_labels)
    )
    labels = [label.strip() for label in custom_labels.split('\n') if label.strip()]
    
    if len(labels) < 2:
        st.error("Please enter at least 2 labels")

# Add a clear-cache button in sidebar
if st.sidebar.button("Clear Cache"):
    st.cache_data.clear()
    st.success("Cache cleared!")

if st.button("Classify"):
    if model_ready and user_input and len(labels) >= 2:
        payload = {
            "inputs": user_input,
            "parameters": {"candidate_labels": labels}
        }
        
        with st.spinner('Classifying...'):
            result = query(payload)
        
        if result and not "error" in result:
            # Display results
            st.subheader("Classification Results:")
            
            # Create a DataFrame for display
            results_df = pd.DataFrame({
                'Label': result['labels'],
                'Score': result['scores']
            })
            
            # Sort by score descending
            results_df = results_df.sort_values('Score', ascending=False)
            
            # Display as table
            st.dataframe(results_df.style.format({'Score': '{:.4f}'}))
            
            # Display as bar chart
            # st.bar_chart(results_df.set_index('Label'))
            
            # Show cache info
            st.sidebar.info("Using cached model")
        else:
            st.error("Failed to get classification results")
    elif len(labels) < 2:
        st.error("Please enter at least 2 labels for classification")

# Display current labels and cache status in sidebar
st.sidebar.subheader("Current Labels:")
st.sidebar.write(labels)
st.sidebar.subheader("Cache Status:")
st.sidebar.write("Model Initialization: Cached" if model_ready else "Model Initialization: Not Cached")