Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, BitsAndBytesConfig, GenerationConfig
 import gradio as gr
 import torch
 
@@ -16,14 +16,8 @@ print("********************")
 print(device)
 print("********************")
 
-bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type='nf4',
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_compute_dtype=torch.bfloat16
-    )
 
-model = AutoModelForCausalLM.from_pretrained(model_id, config=config, quantization_config=bnb_config)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 model = model.to(device)
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)