Update app.py

app.py

@@ -3,11 +3,12 @@ import logging
 import json
 import os
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import psutil
 import cachetools
 import hashlib
+from bitsandbytes import quantize
 
 # Set environment variable for cache
 os.environ["HF_HOME"] = "/app/cache"
@@ -24,13 +25,13 @@ try:
     tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", cache_dir="/app/cache")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
         device_map="cpu",
         low_cpu_mem_usage=True,
         cache_dir="/app/cache",
-        trust_remote_code=True
+        trust_remote_code=True,
+        quantization_config={"load_in_4bit": True}  # 4-bit quantization
     )
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
 except Exception as e:
     logger.error(f"Failed to load BitNet model: {str(e)}")
     raise HTTPException(status_code=500, detail=f"BitNet model initialization failed: {str(e)}")
@@ -48,7 +49,7 @@ def get_text_hash(text: str):
     """Generate MD5 hash of text."""
     return hashlib.md5(text.encode('utf-8')).hexdigest()
 
-# Allowed categories and subcategories
+# Simplified categories (reference only, not included in prompt)
 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"]},
@@ -87,29 +88,29 @@ async def categorize_with_bitnet(description: str, amount: float):
         return structured_data_cache[text_hash]
 
     try:
-        prompt = f"""You are an expert financial transaction categorizer using BitNet b1.2-3B. Given a transaction description and amount, categorize it into the specified categories and subcategories. Assign confidence scores (0 to 1). Follow these rules:
-- Select category and subcategory from this list (case-insensitive, use exact names):
-{', '.join([f'{c["name"]} ({", ".join(c["subcategories"])})' for c in ALLOWED_CATEGORIES])}
-- For positive amounts, use 'income' and one of its subcategories.
-- If unsure, set confidence to 0.7.
-- If no match, use 'miscellaneous' and 'other'.
-- Do NOT add markdown or explanations, only output valid JSON.
+        # Simplified prompt
+        prompt = f"""Categorize this transaction into a category and subcategory with confidence scores (0 to 1). Use 'income' for positive amounts. If unsure, use confidence 0.7 and 'miscellaneous'/'other' if no match. Output only JSON.
 
 Description: {description}
 Amount: {amount}
 
-Output JSON:
 {{
   "category": "",
   "subcategory": "",
   "category_confidence": 0.0,
   "subcategory_confidence": 0.0
-}}
-"""
-        outputs = pipe(prompt)[0]["generated_text"]
-        json_start = outputs.rfind("{")
-        json_end = outputs.rfind("}") + 1
-        result = json.loads(outputs[json_start:json_end])
+}}"""
+        inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=100,  # Reduced for speed
+            do_sample=False,
+            num_beams=1
+        )
+        json_str = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        json_start = json_str.rfind("{")
+        json_end = json_str.rfind("}") + 1
+        result = json.loads(json_str[json_start:json_end])
 
         # Normalize category and subcategory
         def normalize(s):