Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+from peft import prepare_model_for_kbit_training
 from transformers import AutoTokenizer, AutoModelForCausalLM
 model_path = "inception-mbzuai/jais-13b-chat"
 
@@ -10,7 +11,10 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 #model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, trust_remote_code=True)
+#model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, trust_remote_code=True)
+
+model = AutoModelForCausalLM.from_pretrained("inception-mbzuai/jais-13b-chat", load_in_8bit=True, device_map="auto", trust_remote_code=True)
+model = prepare_model_for_kbit_training(model)
 
 def get_response(text,tokenizer=tokenizer,model=model):
     input_ids = tokenizer(text, return_tensors="pt").input_ids