Update app.py

app.py

@@ -132,10 +132,10 @@ class StopOnTokens(StoppingCriteria):
 
 def initialize_model():
     quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.bfloat16,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_use_double_quant=True,
+        load_in_8bit=True,
+        bnb_8bit_compute_dtype=torch.bfloat16,
+        bnb_8bit_quant_type="nf4",
+        bnb_8bit_use_double_quant=True,
     )
 
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
@@ -147,14 +147,14 @@ def initialize_model():
         quantization_config=quantization_config,
         torch_dtype=torch.bfloat16,
         trust_remote_code=True
-    ).to("cuda")
+    )
 
     return model, tokenizer
 
 def format_response(text):
     return text.replace("[Understand]", '\n<strong class="special-tag">[Understand]</strong>\n') \
-              .replace("[Plan]", '\n<strong class="special-tag">[Plan]</strong>\n') \
-              .replace("[Conclude]", '\n<strong class="special-tag">[Conclude]</strong>\n') \
+              .replace("[/Reason]", '\n<strong class="special-tag">[/Reason]</strong>\n') \
+              .replace("[/Answer]", '\n<strong class="special-tag">[/Answer]</strong>\n') \
               .replace("[Reason]", '\n<strong class="special-tag">[Reason]</strong>\n') \
               .replace("[Answer]", '\n<strong class="special-tag">[Answer]</strong>\n')
 @spaces.GPU(duration=360)