Update app.py

app.py

@@ -17,7 +17,7 @@ MODEL_CONFIG = {
     }
 }
 
-# Quantization config for 4-bit loading
+# Quantization config (4-bit)
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -39,9 +39,9 @@ class ChatModel:
             
             model = AutoModelForCausalLM.from_pretrained(
                 config["model_name"],
+                quantization_config=bnb_config,
                 device_map="auto",
                 torch_dtype=torch.float16,
-                attn_implementation="flash_attention_2" if "phi-3" in model_name else "eager",
                 low_cpu_mem_usage=True
             )
             
@@ -60,7 +60,7 @@ class ChatModel:
             "text-generation",
             model=self.models[model_name],
             tokenizer=self.tokenizers[model_name],
-            max_new_tokens=512,
+            max_new_tokens=384,
             temperature=0.7,
             top_p=0.9,
             repetition_penalty=1.1,