Update app.py

app.py

@@ -8,7 +8,6 @@ import torch
 import psutil
 import cachetools
 import hashlib
-from bitsandbytes import quantize
 
 # Set environment variable for cache
 os.environ["HF_HOME"] = "/app/cache"
@@ -25,12 +24,11 @@ try:
     tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", cache_dir="/app/cache")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.bfloat16,  # Optimized for CPU
         device_map="cpu",
         low_cpu_mem_usage=True,
         cache_dir="/app/cache",
-        trust_remote_code=True,
-        quantization_config={"load_in_4bit": True}  # 4-bit quantization
+        trust_remote_code=True
     )
 except Exception as e:
     logger.error(f"Failed to load BitNet model: {str(e)}")
@@ -49,7 +47,7 @@ def get_text_hash(text: str):
     """Generate MD5 hash of text."""
     return hashlib.md5(text.encode('utf-8')).hexdigest()
 
-# Simplified categories (reference only, not included in prompt)
+# Simplified categories (reference only, not 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"]},
@@ -103,7 +101,7 @@ Amount: {amount}
         inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
         outputs = model.generate(
             **inputs,
-            max_new_tokens=100,  # Reduced for speed
+            max_new_tokens=50,  # Further reduced for speed
             do_sample=False,
             num_beams=1
         )