app.py

import gradio as gr
from transformers import pipeline
from typing import Dict, List, Union
import logging
from pydantic import BaseModel, Field

# Logging configuration
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants
MAX_LENGTH = 512
TOP_K = 16
MODELS = {
    "bert": "albertmartinez/bert-sdg-classification",
    "bert-multilingual": "albertmartinez/bert-multilingual-sdg-classification",
    "distilbert-multilingual": "albertmartinez/distilbert-multilingual-sdg-classification",
    "xlm-roberta-large": "albertmartinez/xlm-roberta-large-sdg-classification"
}

class ClassificationResult(BaseModel):
    """Model for classification results."""
    label: str = Field(..., description="Classification label")
    score: float = Field(..., description="Classification score")

def load_models() -> Dict[str, pipeline]:
    """Load all required models."""
    try:
        return {
            name: pipeline("text-classification", model=model_path)
            for name, model_path in MODELS.items()
        }
    except Exception as e:
        logger.error(f"Error loading models: {str(e)}")
        raise

# Load models
models = load_models()

def validate_input(text: str) -> bool:
    """Validate input text."""
    if not text or not isinstance(text, str):
        return False
    if len(text.strip()) == 0:
        return False
    return True

def classify_text(text: str, model: pipeline) -> Dict[str, float]:
    """
    Classify text using the specified model.
    
    Args:
        text: Text to classify
        model: Classification model
        
    Returns:
        Dictionary with labels and their scores
    """
    if not validate_input(text):
        raise ValueError("Invalid text input")
    
    try:
        result = model(text, top_k=TOP_K, truncation=True, max_length=MAX_LENGTH)
        return {p["label"]: p["score"] for p in result}
    except Exception as e:
        logger.error(f"Error in classification: {str(e)}")
        raise

def classify_all(text: str) -> List[Dict[str, float]]:
    """
    Classify text using all available models.
    
    Args:
        text: Text to classify
        
    Returns:
        List of dictionaries with predictions from each model
    """
    if not validate_input(text):
        raise ValueError("Invalid text input")
    
    try:
        return [
            classify_text(text, model)
            for model in models.values()
        ]
    except Exception as e:
        logger.error(f"Error in multiple classification: {str(e)}")
        raise

def create_interface(model_name: str, model: pipeline) -> gr.Interface:
    """Create a Gradio interface for a specific model."""
    return gr.Interface(
        fn=lambda text: classify_text(text, model),
        inputs=gr.Textbox(
            lines=2,
            label="Text",
            placeholder="Enter text here..."
        ),
        outputs=gr.Label(label="Top SDG Predicted"),
        title=f"{model_name.upper()} SDG classification",
        description="Enter text and see the classification results!",
        examples=[["To respond to these new challenges PSV has been encouraged to develop strong links with business and enterprises, to cover remote areas and to cater for less affluent students. It analyses the impact of the dual sector universities in Australia, short cycle higher education in Scotland, three sub-sectors of Norway's tertiary education and the emerging non-university higher education in Italy, as well as the vocational education and training in Spain. It discusses the issues of transition, participation and collaboration of different types of post-secondary education and analysis the impact of policy changes."]],
        flagging_mode="never",
        api_name=f"classify_{model_name}"
    )

# Create main interface
ifaceall = gr.Interface(
    fn=classify_all,
    inputs=gr.Textbox(
        lines=2,
        label="Text",
        placeholder="Enter text here..."
    ),
    outputs=[
        gr.Label(label="bert"),
        gr.Label(label="bert-multilingual"),
        gr.Label(label="distilbert-multilingual"),
        gr.Label(label="xlm-roberta-large")
    ],
    title="SDG text classification",
    description="Enter text and see the classification results!",
    examples=[["To respond to these new challenges PSV has been encouraged to develop strong links with business and enterprises, to cover remote areas and to cater for less affluent students. It analyses the impact of the dual sector universities in Australia, short cycle higher education in Scotland, three sub-sectors of Norway's tertiary education and the emerging non-university higher education in Italy, as well as the vocational education and training in Spain. It discusses the issues of transition, participation and collaboration of different types of post-secondary education and analysis the impact of policy changes."]],
    flagging_mode="never",
    api_name="classify_all_models"
)

# Create individual interfaces
interfaces = [ifaceall] + [
    create_interface(name, model)
    for name, model in models.items()
]

# Create main interface with tabs
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.TabbedInterface(
        interface_list=interfaces,
        tab_names=["ALL"] + list(MODELS.keys()),
        title="Sustainable Development Goals (SDG) Text Classifier App"
    )

if __name__ == "__main__":
    logger.info(f"Gradio version: {gr.__version__}")
    demo.launch()