File size: 1,400 Bytes
0da497e
dd4466b
 
 
0da497e
 
 
dd4466b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
from django.apps import AppConfig
import torch
from transformers import pipeline
import os


class ApiConfig(AppConfig):
    default_auto_field = "django.db.models.BigAutoField"
    name = "api"

    def ready(self):
        # Prevent loading twice (Django sometimes runs ready() twice in dev mode)
        if os.environ.get("RUN_MAIN", None) != "true" and not os.environ.get(
            "KUBERNETES_PORT"
        ):
            pass  # Keep going if running via Gunicorn or in production

        print("Loading Expense Categorizer AI from cache...")

        device = 0 if torch.cuda.is_available() else -1

        # This will automatically use the cached version built by Docker
        self.classifier = pipeline(
            "zero-shot-classification",
            model="valhalla/distilbart-mnli-12-3",
            device=device,
        )

        self.categories = [
            "Food & Drinks",
            "Groceries",
            "Shopping",
            "Bills & Utilities",
            "Entertainment",
            "Health",
            "Education",
            "Subscriptions",
            "Travel",
            "Rent",
            "Family & Friends",
            "Miscellaneous",
            "Gifts",
            "Party",
            "Personal Care",
            "Home & Hygiene",
            "Others",
            "Recharge",
        ]
        print("AI Model loaded successfully!")