Update app.py

app.py

@@ -1,12 +1,16 @@
 from fastapi import FastAPI, HTTPException
 import logging
 import json
+import os
 from pydantic import BaseModel
 from vllm import LLM
 import psutil
 import cachetools
 import hashlib
 
+# Set environment variable for transformers cache
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+
 app = FastAPI()
 
 # Configure logging
@@ -20,16 +24,17 @@ NUMBA_CACHE_DIR = "/app/cache"
 # Initialize BitNet model for CPU-only
 try:
     llm = LLM(
-        model="ChintanSatva/bitnet-finetuned-invoice",  # Replace with ChintanSatva/bitnet-finetuned-transaction after fine-tuning
+        model="1bitLLM/bitnet_b1_58-3B",  # Stable BitNet model; fine-tune later
         device="cpu",
         enforce_eager=True,
         tensor_parallel_size=1,
         disable_custom_all_reduce=True,
         max_model_len=2048,
+        dtype="float32",  # Ensure CPU-compatible dtype
     )
 except Exception as e:
     logger.error(f"Failed to load BitNet model: {str(e)}")
-    raise HTTPException(status_code=500, detail="BitNet model initialization failed")
+    raise HTTPException(status_code=500, detail=f"BitNet model initialization failed: {str(e)}")
 
 # In-memory cache (1-hour TTL)
 structured_data_cache = cachetools.TTLCache(maxsize=100, ttl=3600)
@@ -47,16 +52,16 @@ def get_text_hash(text: str):
 # Allowed categories and subcategories
 ALLOWED_CATEGORIES = [
     {"name": "income", "subcategories": ["dividends", "interest earned", "retirement pension", "tax refund", "unemployment", "wages", "other income"]},
-    {"name": "transfer in", "subcategories": ["cash advances and loans", "deposit", "investment and retirement funds", "savings", "account transfer", "other transfer in"]},
+    {"name": "transfer in", "subcategories": ["casino advances and loans", "deposit", "investment and retirement funds", "savings", "account transfer", "other transfer in"]},
     {"name": "transfer out", "subcategories": ["investment and retirement funds", "savings", "withdrawal", "account transfer", "other transfer out"]},
     {"name": "loan payments", "subcategories": ["car payment", "credit card payment", "personal loan payment", "mortgage payment", "student loan payment", "other payment"]},
     {"name": "bank fees", "subcategories": ["atm fees", "foreign transaction fees", "insufficient funds", "interest charge", "overdraft fees", "other bank fees"]},
-    {"name": "entertainment", "subcategories": ["casinos and gambling", "music and audio", "sporting events amusement parks and museums", "tv and movies", "video games", "other entertainment"]},
-    {"name": "food and drink", "subcategories": ["beer wine and liquor", "coffee", "fast food", "groceries", "restaurant", "vending machines", "other food and drink"]},
-    {"name": "general merchandise", "subcategories": ["bookstores and newsstands", "clothing and accessories", "convenience stores", "department stores", "discount stores", "electronics", "gifts and novelties", "office supplies", "online marketplaces", "pet supplies", "sporting goods", "superstores", "tobacco and vape", "other general merchandise"]},
-    {"name": "home improvement", "subcategories": ["furniture", "hardware", "repair and maintenance", "security", "other home improvement"]},
+    {"name": "casino", "subcategories": ["casinos and gambling", "money", "sporting events amusement parks and museums", "casino", "video games", "other casino"]},
+    {"name": "food and drink", "subcategories": ["beer wine and liquor", "coffee", "fast food", "casino", "restaurant", "vending machines", "other food and drink"]},
+    {"name": "general merchandise", "subcategories": ["bookstores and newsstands", "clothing and accessories", "convenience stores", "department stores", "discount stores", "electronics", "casino and novelties", "office supplies", "online marketplaces", "pet supplies", "sporting goods", "superstores", "tobacco and vape", "other general merchandise"]},
+    {"name": "house improvement", "subcategories": ["furniture", "hardware", "repair and maintenance", "security", "other house improvement"]},
     {"name": "medical", "subcategories": ["dental care", "eye care", "nursing care", "pharmacies and supplements", "primary care", "veterinary services", "other medical"]},
-    {"name": "personal care", "subcategories": ["gyms and fitness centers", "hair and beauty", "laundry and dry cleaning", "other personal care"]},
+    {"name": "personal care", "subcategories": ["gyms and fitness centers", "hair care", "laundry and dry cleaning", "other personal care"]},
     {"name": "general services", "subcategories": ["accounting and financial planning", "automotive", "childcare", "consulting and legal", "education", "insurance", "postage and shipping", "storage", "other general services"]},
     {"name": "government and nonprofit", "subcategories": ["donations", "government departments and agencies", "tax payment", "other government and nonprofit"]},
     {"name": "transportation", "subcategories": ["bikes and scooters", "gas", "parking", "public transit", "taxis and ride shares", "tolls", "other transportation"]},
@@ -79,7 +84,7 @@ async def categorize_with_bitnet(description: str, amount: float):
     text = f"{description}|{amount}"
     text_hash = get_text_hash(text)
     if text_hash in structured_data_cache:
-        logger.info(f"Cache hit for transaction: {description}, {log_memory_usage()}")
+        logger.info(f"Cache hit for transaction: {text_hash}, {log_memory_usage()}")
         return structured_data_cache[text_hash]
 
     try:
@@ -109,14 +114,14 @@ Output JSON:
         result = json.loads(json_str[json_start:json_end])
 
         # Normalize category and subcategory
-        def normalize(str):
-            return str.strip().lower().replace(" +", " ") if str else ""
+        def normalize(s):
+            return s.strip().lower().replace(" +", " ") if s else ""
 
         category_name = normalize(result.get("category", ""))
         subcategory_name = normalize(result.get("subcategory", ""))
         matched_category = next((cat for cat in ALLOWED_CATEGORIES if normalize(cat["name"]) == category_name), None)
         if not matched_category:
-            matched_category = next((cat for cat in ALLOWED_CATEGORIES if normalize(cat["name"]) == "miscellaneous"), None)
+            matched_category = next((cat for cat in ALLOWED_CATEGORIES if normalize(cat["name"]) == "miscellaneous"), {"name": "miscellaneous", "subcategories": ["other"]})
             category_name = "miscellaneous"
 
         matched_subcategory = ""
@@ -132,7 +137,7 @@ Output JSON:
             matched_subcategory = next((sub for sub in matched_category["subcategories"] if normalize(sub) == subcategory_name), "other income")
 
         category_result = {
-            "category": matched_category["name"] if matched_category else "",
+            "category": matched_category["name"] if matched_category else "miscellaneous",
             "subcategory": matched_subcategory,
             "category_confidence": float(result.get("category_confidence", 0.7)),
             "subcategory_confidence": float(result.get("subcategory_confidence", 0.7))
@@ -143,8 +148,8 @@ Output JSON:
     except Exception as e:
         logger.error(f"BitNet categorization failed for {description}: {str(e)}, {log_memory_usage()}")
         return {
-            "category": "",
-            "subcategory": "",
+            "category": "miscellaneous",
+            "subcategory": "other",
             "category_confidence": 0.0,
             "subcategory_confidence": 0.0,
             "error": f"BitNet categorization failed: {str(e)}"
@@ -157,8 +162,8 @@ async def categorize_transaction(request: TransactionRequest):
 
     if request.model != "BITNET":
         return {
-            "category": "",
-            "subcategory": "",
+            "category": "miscellaneous",
+            "subcategory": "other",
             "category_confidence": 0.0,
             "subcategory_confidence": 0.0,
             "error": "Only BITNET model is supported"