added quantization for performance

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from functools import lru_cache
 import logging
 
@@ -12,6 +12,16 @@ logging.basicConfig(level=logging.INFO)
 logger.info("Starting Jerome Powell AI Assistant...")
 
 
+QUANTIZATION_CONFIG = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype="float16",
+)
+
+MODEL_NAME = "BoostedJonP/powell-phi3-mini"
+
+
 @lru_cache(maxsize=1)
 def load_model():
     """Load the fine-tuned Jerome Powell model"""
@@ -26,7 +36,7 @@ def load_model():
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
             trust_remote_code=True,
-            torch_dtype=torch.float16,
+            quantization_config=QUANTIZATION_CONFIG,
             device_map="auto",
             attn_implementation="eager",
             use_cache=True,