Update app.py

app.py

@@ -1,13 +1,30 @@
 import spaces
 import gradio as gr
 import torch
-from transformers import pipeline
+from transformers import pipeline, BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
+
+MODEL_ID = "marcelbinz/Llama-3.1-Centaur-70B"
+
+bnb_4bit_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",
+    attn_implementation="flash_attention_2",
+    quantization_config=bnb_4bit_config,
+)
 
 pipe = pipeline(
     "text-generation",
-    model="marcelbinz/Llama-3.1-Minitaur-8B",
+    model=model,
+    tokenizer=tokenizer,
     device_map="auto",
-    torch_dtype=torch.bfloat16
 )
 
 @spaces.GPU