Update app.py

app.py

@@ -2,31 +2,26 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 import torch
+import html  # เพิ่ม html library
 
 # --- Configuration ---
-BASE_MODEL_ID = "google/gemma-2-2b-it"  # โมเดลหลัก
-ADAPTER_ID = "Phonsiri/gemma-2-2b-it-grpo-v6-checkpoints"  # โมเดล Adapter (LoRA) ที่ต้องการโหลด
+BASE_MODEL_ID = "google/gemma-2-2b-it"
+ADAPTER_ID = "Phonsiri/gemma-2-2b-it-grpo-v6-checkpoints"
 
 # --- Load Tokenizer & Model ---
 print(f"Loading base model: {BASE_MODEL_ID}...")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
 
-# 1. โหลด Base Model ก่อน
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL_ID,
     device_map="auto",
     torch_dtype=torch.float16
 )
 
-# 2. โหลด Adapter (LoRA) มาประกบ
 print(f"Loading adapter: {ADAPTER_ID}...")
 model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
 
-# (Optional) ถ้าต้องการให้ Inference เร็วขึ้นนิดหน่อย สามารถ Merge ได้เลย (กิน RAM ตอนโหลดเพิ่มชั่วคราว)
-# model = model.merge_and_unload() 
-
 def generate(prompt):
-    # สร้าง Chat Template
     messages = [{"role": "user", "content": prompt}]
     
     formatted_prompt = tokenizer.apply_chat_template(
@@ -37,39 +32,35 @@ def generate(prompt):
     
     inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
     
-    # Generate
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=2048,  # เพิ่มความยาวเผื่อการคิดแบบ Chain-of-Thought
-            temperature=0.6,      # ลดลงนิดหน่อยเพื่อให้ Reasoning นิ่งขึ้น
+            max_new_tokens=2048,
+            temperature=0.6,
             top_p=0.9,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
             eos_token_id=tokenizer.eos_token_id
         )
     
-    # Decode
     full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    # --- Response Cleaning Logic ---
-    # ตัดส่วน Prompt ออกเพื่อให้เหลือแค่คำตอบของโมเดล
+    # Cleaning Logic
     if "model\n" in full_response:
-        # ตัดที่ token model ตัวสุดท้าย (Gemma chat format)
         response = full_response.split("model\n")[-1].strip()
     elif "<start_of_turn>model" in full_response:
         response = full_response.split("<start_of_turn>model")[-1].strip()
     else:
-        # Fallback: ตัดตามความยาว prompt
-        # (วิธีนี้อาจไม่แม่นยำ 100% ถ้า prompt ถูก format ใหม่ แต่ใช้กันเหนียว)
-        response = full_response[len(formatted_prompt):].strip() # ตัดจาก formatted prompt ดีกว่า
-        if len(response) == 0: # ถ้าตัดแล้วหายหมด ให้ใช้ raw decode
+        response = full_response[len(formatted_prompt):].strip()
+        if len(response) == 0:
              response = full_response
 
-    # ลบ tags ที่อาจหลงเหลือ
     response = response.replace("<end_of_turn>", "").strip()
     
-    return response
+    # --- สำคัญ: แก้ไขการแสดงผล Tag ---
+    # แปลง < เป็น &lt; เพื่อให้ Gradio ไม่มองว่าเป็น HTML tag ที่ต้องซ่อน
+    # หรือใช้วิธีใส่ Code Block ครอบ
+    return f"```xml\n{response}\n```"
 
 # --- Gradio UI ---
 examples = [
@@ -78,6 +69,7 @@ examples = [
     ["Solve for x: 2x + 5 = 15"]
 ]
 
+# เปลี่ยน Output เป็น Markdown เพื่อให้ render code block สวยๆ
 demo = gr.Interface(
     fn=generate,
     inputs=gr.Textbox(
@@ -85,9 +77,8 @@ demo = gr.Interface(
         lines=3, 
         placeholder="Ask a math or reasoning question..."
     ),
-    outputs=gr.Textbox(
-        label="Reasoning & Answer", 
-        lines=15  # เพิ่มบรรทัดเพราะ GRPO มักตอบยาว
+    outputs=gr.Markdown(  # เปลี่ยนจาก Textbox เป็น Markdown
+        label="Reasoning & Answer"
     ),
     title="Gemma-2-2B GRPO (Adapter Version)",
     description=f"Running Adapter: {ADAPTER_ID}\nBase Model: {BASE_MODEL_ID}",