src/ai_classifier.py

import streamlit as st
from transformers import pipeline

@st.cache_resource
def load_classifier():
    return pipeline(
        "zero-shot-classification",
        model="joeddav/xlm-roberta-large-xnli",  # or use a smaller model
        cache_dir="./hf_cache"
    )

classifier = load_classifier()

CATEGORIES = {
    "Family": "కుటుంబం",
    "Friendship": "స్నేహం",
    "Morality": "నీతి",
    "Hard Work": "శ్రమ",
    "Knowledge": "జ్ఞానం",
    "Devotion": "భక్తి",
    "Culture": "సంస్కృతి",
    "Literature": "సాహిత్యం",
    "Humility": "వినయం",
    "Patience": "సహనం",
    "Courage": "ధైర్యం",
    "Arrogance": "అహంకారం",
    "Love": "ప్రేమ",
    "Greed": "దురాశ",
    "Wisdom": "ఆలోచన",
    "Responsibility": "బాధ్యత",
    "Satire": "వ్యంగ్యం",
    "Politics": "రాజకీయం",
    "Wealth": "ధనము",
    "Time": "సమయం"
}

def classify_proverb(text):
    result = classifier(text, list(CATEGORIES.keys()))
    top_label = result["labels"][0]
    return CATEGORIES[top_label]