src/utils.py

import streamlit as st
from typing import Dict, Any
from transformers import pipeline
from src.constants import SETTINGS, COLORS


@st.cache_resource
def load_models() -> Dict[str, pipeline]:
    return {
        model_name: pipeline(task=SETTINGS.task, model=model_name) 
        for model_name in SETTINGS.models
    }


def colorize_sentiment(label, score) -> str:
    color = COLORS.get(label, COLORS['DEFAULT'])

    return f"""
    <div style="
        display: inline-block;
        background-color: {color};
        color: white;
        padding: 0.2em 0.5em;
        border-radius: 0.3em;
        margin: 0.1em;
        font-size: 0.9em;
        position: relative;
    ">{label}</div>
    <div style="
        font-size: 1em;
        color: {color};
        margin-left: 0.5em
    ">{score:.2f}</div>
    """


def print_results(result: Dict[str, Any]):
    label = result['label']
    score = result['score']
    
    st.markdown(f"""
    <div style="display: flex; align-items: center; margin-bottom: 1em;">
        <span style="font-size: 1.1em; margin-right: 1em;">Overall sentiment:</span>
        {colorize_sentiment(label, score)}
    </div>
    """, unsafe_allow_html=True)