Update app.py

app.py

@@ -15,7 +15,7 @@ from transformers import (
     StoppingCriteriaList
 )
 
-MODEL_ID = "open-thoughts/OpenThinker-32B"
+MODEL_ID = "Daemontatox/mini-Cogito-R1"
 
 DEFAULT_SYSTEM_PROMPT = """
 
@@ -60,7 +60,7 @@ def initialize_model():
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         device_map="cuda",
-        quantization_config=quantization_config,
+        #quantization_config=quantization_config,
         torch_dtype=torch.bfloat16,
         trust_remote_code=True,
         attn_implementation="flash_attention_2"
@@ -97,7 +97,7 @@ def apply_llama3_chat_template(conversation, add_generation_prompt=True):
         elif role == "USER":
             prompt += "<|USER|>\n" + msg["content"].strip() + "\n"
         elif role == "ASSISTANT":
-            prompt += "<|ASSISTANT|>\n" + msg["content"].strip() + "\n"
+            prompt += "<|ASSISTANT|>\n" + msg["content"].strip() + "<think>\n"
     if add_generation_prompt:
         prompt += "<|ASSISTANT|>\n"
     return prompt