Update app.py

app.py

@@ -9,7 +9,8 @@ from threading import Thread
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = "CohereForAI/aya-23-8B"
-MODEL_NAME = MODEL_ID.split("/")[-1]
+MODEL_ID2 = "CohereForAI/aya-23-35B"
+MODEL_NAME = MODEL_ID2.split("/")[-1]
 
 TITLE = "<h1><center>Aya-23-Chatbox</center></h1>"
 
@@ -34,26 +35,27 @@ USE_FLASH_ATTENTION = False
 GRAD_ACC_STEPS = 16
 
 quantization_config = None
+
 if QUANTIZE_4BIT:
-  quantization_config = BitsAndBytesConfig(
-      load_in_4bit=True,
-      bnb_4bit_quant_type="nf4",
-      bnb_4bit_use_double_quant=True,
-      bnb_4bit_compute_dtype=torch.bfloat16,
-  )
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+    )
 
 attn_implementation = None
 if USE_FLASH_ATTENTION:
-  attn_implementation="flash_attention_2"
+    attn_implementation="flash_attention_2"
 
 model = AutoModelForCausalLM.from_pretrained(
-          MODEL_ID,
+          MODEL_ID2,
           quantization_config=quantization_config,
           attn_implementation=attn_implementation,
           torch_dtype=torch.bfloat16,
           device_map="auto",
         )
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID2)
 
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int):