Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 from dotenv import load_dotenv
 from huggingface_hub import login
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 import gradio as gr
 
 load_dotenv()
@@ -12,7 +12,17 @@ login(API_TOKEN)
 
 model_id = "meta-llama/Llama-2-7b-chat-hf"
 
-model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True, device_map="auto")
+quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
+device_map = {
+    "transformer.word_embeddings": 0,
+    "transformer.word_embeddings_layernorm": 0,
+    "lm_head": "cpu",
+    "transformer.h": 0,
+    "transformer.ln_f": 0,
+}
+
+
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device_map,quantization_config=quantization_config)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 generate_text_pipeline = pipeline(