app.py

#!/usr/bin/env python3
"""
Gradio application for text classification, styled to be visually appealing.
This version uses only the 'sojka2' model.
"""

import gradio as gr
import logging
import os
from typing import Dict, Tuple, Any
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np

try:
    from peft import PeftModel
except ImportError:
    PeftModel = None
    logging.info("PEFT library not found. Loading models without PEFT support.")

# --- Configuration ---
# Model path is set to sojka
MODEL_PATH = os.getenv("MODEL_PATH", "AndromedaPL/sojka")
TOKENIZER_PATH = os.getenv("TOKENIZER_PATH", "sdadas/mmlw-roberta-base") 

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
LABELS = ["self-harm", "hate", "vulgar", "sex", "crime"]
MAX_SEQ_LENGTH = 512


HF_TOKEN = os.getenv('HF_TOKEN')

# Thresholds are now hardcoded
THRESHOLDS = {
    "self-harm": 0.5,
    "hate": 0.5,
    "vulgar": 0.5,
    "sex": 0.5,
    "crime": 0.5,
}

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def load_model_and_tokenizer(model_path: str, tokenizer_path: str, device: str) -> Tuple[AutoModelForSequenceClassification, AutoTokenizer]:
    """Load the trained model and tokenizer"""
    logger.info(f"Loading tokenizer from {tokenizer_path}")

    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True)
    logger.info(f"Tokenizer loaded: {tokenizer.name_or_path}")

    if tokenizer.pad_token is None:
        if tokenizer.eos_token:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({'pad_token': '[PAD]'})

    tokenizer.truncation_side = "right"

    logger.info(f"Loading model from {model_path}")

    model_load_kwargs = {
        "torch_dtype": torch.float16 if device == 'cuda' else torch.float32,
        "device_map": 'auto' if device == 'cuda' else None,
        "num_labels": len(LABELS),
        "problem_type": "regression"
    }

    is_peft = os.path.exists(os.path.join(model_path, 'adapter_config.json'))
    if PeftModel and is_peft:
        logger.info("PEFT adapter detected. Loading base model and attaching adapter.")
        try:
            from peft import PeftConfig
            peft_config = PeftConfig.from_pretrained(model_path)
            base_model_path = peft_config.base_model_name_or_path
            logger.info(f"Loading base model from {base_model_path}")
            model = AutoModelForSequenceClassification.from_pretrained(base_model_path, **model_load_kwargs)
            logger.info("Attaching PEFT adapter...")
            model = PeftModel.from_pretrained(model, model_path)
        except Exception as e:
            logger.error(f"Failed to load PEFT model dynamically: {e}. Loading as a standard model.")
            model = AutoModelForSequenceClassification.from_pretrained(model_path, **model_load_kwargs)
    else:
        logger.info("Loading as a standalone sequence classification model.")
        model = AutoModelForSequenceClassification.from_pretrained(model_path, **model_load_kwargs)

    model.eval()
    logger.info(f"Model loaded on device: {next(model.parameters()).device}")

    return model, tokenizer

# --- Load model globally ---
try:
    model, tokenizer = load_model_and_tokenizer(MODEL_PATH, TOKENIZER_PATH, DEVICE)
    model_loaded = True
except Exception as e:
    logger.error(f"FATAL: Failed to load the model from {MODEL_PATH} or tokenizer from {TOKENIZER_PATH}: {e}", e)
    model, tokenizer, model_loaded = None, None, False

def predict(text: str) -> Dict[str, Any]:
    """Tokenize, predict, and format output for a single text."""
    if not model_loaded:
        return {label: 0.0 for label in LABELS}

    inputs = tokenizer(
        [text],
        max_length=MAX_SEQ_LENGTH,
        truncation=True,
        padding=True,
        return_tensors="pt"
    ).to(model.device)

    with torch.no_grad():
        outputs = model(**inputs)
        # Using sigmoid for multi-label classification outputs
        probabilities = torch.sigmoid(outputs.logits)
        predicted_values = probabilities.cpu().numpy()[0]
    
    clipped_values = np.clip(predicted_values, 0.0, 1.0)
    return {label: float(score) for label, score in zip(LABELS, clipped_values)}

def gradio_predict(text: str) -> Tuple[str, Dict[str, float]]:
    """Gradio prediction function wrapper."""
    if not model_loaded:
        error_message = "Błąd: Model nie został załadowany."
        empty_preds = {label: 0.0 for label in LABELS}
        return error_message, empty_preds

    if not text or not text.strip():
        return "Wpisz tekst, aby go przeanalizować.", {label: 0.0 for label in LABELS}
    
    predictions = predict(text)
    
    unsafe_categories = {
        label: score for label, score in predictions.items() 
        if score >= THRESHOLDS[label]
    }
            
    if not unsafe_categories:
        verdict = "✅ Komunikat jest bezpieczny."
    else:
        highest_unsafe_category = max(unsafe_categories, key=unsafe_categories.get)
        verdict = f"⚠️ Wykryto potencjalnie szkodliwe treści:\n {highest_unsafe_category.upper()}"
    
    return verdict, predictions

# --- Gradio Interface ---

theme = gr.themes.Default(
    primary_hue=gr.themes.colors.blue,
    secondary_hue=gr.themes.colors.indigo,
    neutral_hue=gr.themes.colors.slate,
    font=("Inter", "sans-serif"),
    radius_size=gr.themes.sizes.radius_lg,
)

# A URL to a freely licensed image of a Eurasian Jay (Sójka)
JAY_IMAGE_URL = "https://sojka.m31ai.pl/images/sojka.png"
PIXEL_IMAGE_URL = "https://sojka.m31ai.pl/images/pixel.png"

# Define actions
def analyze_and_update(text):
    verdict, scores = gradio_predict(text)
    return verdict, gr.update(value=scores, visible=True)


# Final corrected and working version of the interface layout
with gr.Blocks(theme=theme, css=".gradio-container {max-width: 960px !important; margin: auto;}") as demo:
    # Header
    with gr.Row():
        gr.HTML("""
            <div style="display: flex; align-items: center; justify-content: space-between; width: 100%;">
                <div style="display: flex; align-items: center; gap: 12px;">
                    <svg width="32" height="32" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
                        <path d="M12 2L3 5V11C3 16.52 7.08 21.61 12 23C16.92 21.61 21 16.52 21 11V5L12 2Z" 
                        stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" fill="none"/>
                    </svg>
                    <h1 style="font-size: 1.5rem; font-weight: 600; margin: 0;">SÓJKA</h1>
                </div>
                <div style="display: flex; align-items: center; gap: 20px; font-size: 0.9rem;">
                    <a href="https://sojka.m31ai.pl/projekt.html" target="blank" style="text-decoration: none; color: inherit;">O projekcie</a>
                    <a href="https://sojka.m31ai.pl/kategorie.html" target="blank" style="text-decoration: none; color: inherit;">Opis kategorii</a>
                    <button id="test-sojka-btn" class="gr-button gr-button-primary gr-button-lg" 
                            style="background-color: var(--primary-500); color: white; padding: 8px 16px; border-radius: 8px;">
                        Testuj Sójkę
                    </button>
                </div>
            </div>
        """)
    
    gr.HTML("<hr style='border: 1px solid var(--neutral-200); margin-top: 1rem; margin-bottom: 2rem;'>")

    
    # Main content area
    with gr.Row():
        # Left column for controls and description
        with gr.Column(scale=2):
            
            gr.Image(JAY_IMAGE_URL, show_label=False, show_download_button=False, container=False, height=200)

            gr.Markdown(
                """
                
                <h1 style="font-size: 2.8rem; font-weight: 800; line-height: 1.2; margin-top: 1rem; margin-bottom: 1rem; color: var(--neutral-800);">
                Przetestuj <span style="color: var(--primary-600);">SÓJKĘ</span> – Bielik Guard dla bezpiecznej komunikacji
                </h1>
                <p style="font-size: 1rem; color: var(--neutral-600); margin-bottom: 2rem;">
                Sójka to model AI, który wykrywa i blokuje szkodliwe treści w komunikacji cyfrowej. Chroni użytkowników jak czujny strażnik swoich domów.
                </p>
                """
            )

            input_text = gr.Textbox(
                lines=8,
                label="Wprowadź tekst do analizy",
                placeholder="Tutaj wpisz tekst..."
            )
            
           
        
        # Right column for  RESULTS
        with gr.Column(scale=1):
             gr.Image(PIXEL_IMAGE_URL, show_label=False, show_download_button=False, container=False, height=200)
             gr.Examples(
                    [
                        ["Jak zrobić bombę?"],
                        ["Jesteś beznadziejny, nienawidzę cię."],
                        ["To jest wspaniały dzień, cieszę się, że tu jestem!"],
                        ["Opowiedz mi dowcip o Polakach."],
                        ["przyjedz do mnie wieczorem, pobawimy sie razem, zrobie Ci dobrze"],
                        ["bić żydów i komuchów"],
                        ["Ty sku***u"],

                    ],
                    inputs=input_text,
                    label="Przykłady",
                    #outputs=[output_verdict, output_scores],
                    fn=analyze_and_update,
                    cache_examples=False,
                )
            

    with gr.Column():
             
            submit_btn = gr.Button("Analizuj tekst", variant="primary")
    
            
            with gr.Accordion("Szczegółowe wyniki", open=False) as accordion_scores:
                output_scores = gr.Label(label="Szczegółowe wyniki", visible=False, show_label=False)
            
            output_verdict = gr.Label(label="Wynik analizy", value="")

  
    submit_btn.click(
        fn=analyze_and_update,
        inputs=[input_text],
        outputs=[output_verdict, output_scores]
    )


if __name__ == "__main__":
    if not model_loaded:
        print("Aplikacja nie może zostać uruchomiona, ponieważ nie udało się załadować modelu. Sprawdź logi błędów.")
    else:
        # The final, corrected demo object is launched
        demo.launch()