Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
 from transformers import AutoTokenizer, GPTJForCausalLM
+import torch
 
 model_name = "rycont/kakaobrain__kogpt-6b-8bit"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = GPTJForCausalLM.from_pretrained(model_name)
+model = GPTJForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
 
 def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
     outputs = model.generate(inputs['input_ids'], max_new_tokens=50)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
@@ -15,7 +16,7 @@ iface = gr.Interface(
     fn=generate_response,
     inputs="text",
     outputs="text",
-    title="KoGPT-6B Chatbot",
+    title="KoGPT-6B 8bit Chatbot",
     description="Enter a prompt and the model will generate a response."
 )