Update app.py

app.py

@@ -1,10 +1,14 @@
 import gradio as gr
-from transformers import AutoTokenizer, GPTJForCausalLM
+from transformers import AutoTokenizer, GPTJForCausalLM, BitsAndBytesConfig
 import torch
 
 model_name = "rycont/kakaobrain__kogpt-6b-8bit"
+
+# BitsAndBytesConfig 설정
+bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = GPTJForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+model = GPTJForCausalLM.from_pretrained(model_name, device_map="auto", quantization_config=bnb_config)
 
 def generate_response(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")