Update app.py

app.py

@@ -8,13 +8,13 @@ from peft import PeftModel
 # CONFIG
 # =========================
 BASE_MODEL = "google/gemma-3-270m-it"
-LORA_MODEL = "loboGOAT/DogeAI-v1.0-instruct"
+LORA_MODEL = "AxionLab-official/DogeAI-v1.0-instruct"
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # =========================
-# LOAD MODEL
+# LOAD TOKENIZER
 # =========================
 print("🔄 Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(
@@ -22,6 +22,9 @@ tokenizer = AutoTokenizer.from_pretrained(
     token=HF_TOKEN
 )
 
+# =========================
+# LOAD BASE MODEL
+# =========================
 print("🔄 Loading base model...")
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
@@ -30,14 +33,37 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="cpu"
 )
 
+# =========================
+# APPLY LORA
+# =========================
 print("🔄 Applying LoRA...")
 model = PeftModel.from_pretrained(model, LORA_MODEL)
 model.eval()
 
+# =========================
+# OUTPUT PARSER
+# =========================
+def extract_answer(text, show_reasoning=False):
+    if "<think>" in text and "</think>" in text:
+        reasoning = text.split("<think>")[1].split("</think>")[0].strip()
+        answer = text.split("</think>")[-1].strip()
+
+        if show_reasoning:
+            return (
+                "🧠 RACIOCÍNIO INTERNO:\n"
+                f"{reasoning}\n\n"
+                "✅ RESPOSTA FINAL:\n"
+                f"{answer}"
+            )
+        else:
+            return answer
+
+    return text.strip()
+
 # =========================
 # CHAT FUNCTION
 # =========================
-def chat(user_input, system_prompt, temperature, top_p, max_tokens):
+def chat(user_input, system_prompt, temperature, top_p, max_tokens, show_reasoning):
     prompt = (
         "<bos>\n"
         "<start_of_turn>system\n"
@@ -45,6 +71,9 @@ def chat(user_input, system_prompt, temperature, top_p, max_tokens):
         "<start_of_turn>user\n"
         f"{user_input}\n"
         "<start_of_turn>model\n"
+        "<think>\n"
+        "Explique passo a passo seu raciocínio antes de responder.\n"
+        "</think>\n"
     )
 
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
@@ -64,9 +93,9 @@ def chat(user_input, system_prompt, temperature, top_p, max_tokens):
     text = tokenizer.decode(output[0], skip_special_tokens=True)
 
     if "<start_of_turn>model" in text:
-        return text.split("<start_of_turn>model")[-1].strip()
+        text = text.split("<start_of_turn>model")[-1]
 
-    return text.strip()
+    return extract_answer(text, show_reasoning)
 
 # =========================
 # GRADIO UI
@@ -74,13 +103,13 @@ def chat(user_input, system_prompt, temperature, top_p, max_tokens):
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🐕 DogeAI v1.0  
+        # 🐕 DogeAI v1.0 — Reasoning Mode  
 
         ⚠️ **AVISO IMPORTANTE**  
-        Este modelo é **experimental**, pequeno e pode **ALUCINAR** respostas,  
-        **inventar fatos**, errar datas, nomes ou explicações técnicas.
+        Este modelo é **experimental**, pequeno e pode **alucinar**,  
+        inventar fatos ou cometer erros.
 
-        ❌ **NÃO use como fonte confiável**  
+        ❌ Não use como fonte confiável  
         ✅ Use para estudo, testes e experimentação
         """
     )
@@ -90,13 +119,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             user_input = gr.Textbox(
                 lines=5,
                 label="Mensagem",
-                placeholder="Fale com o DogeAI 🐶"
+                placeholder="Converse com o DogeAI 🐶"
             )
 
             submit = gr.Button("Enviar 🚀")
 
             output = gr.Textbox(
-                lines=10,
+                lines=14,
                 label="Resposta do modelo"
             )
 
@@ -112,12 +141,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Instruções internas"
             )
 
-            gr.Markdown("### ⚠️ Hiperparâmetros")
-            gr.Markdown(
-                "**SÓ MEXA SE SOUBER O QUE ESTÁ FAZENDO**/**ONLY ADJUST VALUES IF YOU KNOW WHAT YOU'RE DOING**\n\n"
-                "Valores errados podem causar respostas ruins ou nonsense./Bad values can cause trash or nonsense responses"
+            show_reasoning = gr.Checkbox(
+                value=False,
+                label="Mostrar raciocínio interno (thinking)"
             )
 
+            gr.Markdown("### ⚙️ Hiperparâmetros")
             temperature = gr.Slider(
                 0.1, 1.5, value=0.65, step=0.05, label="Temperature"
             )
@@ -132,9 +161,17 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     submit.click(
         chat,
-        inputs=[user_input, system_prompt, temperature, top_p, max_tokens],
+        inputs=[
+            user_input,
+            system_prompt,
+            temperature,
+            top_p,
+            max_tokens,
+            show_reasoning
+        ],
         outputs=output
     )
 
 demo.launch()
 
+