Update app.py

app.py

@@ -3,9 +3,17 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 # Load the model and tokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import bitsandbytes as bnb
+
 model_name = "microsoft/phi-2"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    load_in_8bit=True,  
+    device_map="auto"
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")