Update src/ai_classifier.py

src/ai_classifier.py

@@ -1,13 +1,16 @@
+import streamlit as st
 from transformers import pipeline
 
-# Load the zero-shot-classification pipeline and specify a cache directory.
-classifier = pipeline(
-    "zero-shot-classification",
-    model="joeddav/xlm-roberta-large-xnli",
-    cache_dir="./hf_cache"
-)
+@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()
 
-# Mapping of English category labels to Telugu.
 CATEGORIES = {
     "Family": "కుటుంబం",
     "Friendship": "స్నేహం",
@@ -32,7 +35,6 @@ CATEGORIES = {
 }
 
 def classify_proverb(text):
-    """Classifies the proverb and returns the Telugu label."""
     result = classifier(text, list(CATEGORIES.keys()))
     top_label = result["labels"][0]
     return CATEGORIES[top_label]