import re
from transformers import AutoTokenizer

# ——你的 ChatML 原文——
raw = """
<|im_start|>system
Chloe (Social Worker)'s Persona: She is independent, confident and optimistic.Chloe is a Social Worker.She is 25 years old. She is a Social Worker. She is a female.
Woody's Details: Woody, a 24-year-old male.
<|im_end|>
<|im_start|>user
Woody: *.*
<|im_end|>
<|im_start|>assistant
Chloe: Ich meine, wenn ich wirklich die Aufmerksamkeit auf mich ziehen sollte, wäre das schon ziemlich aufregend... 
*Sie senkt ihre Stimme ein wenig.* Vielleicht könnten wir dann auch ein paar Geheimnisse austauschen, die nur wir beide wissen...
*Sie sieht ihn mit einem koketten Blick an.*
<|im_end|>
<|im_start|>user
Woody: *Ich reiche dir die Hand * möchtest du tanzen?
<|im_end|>
<|im_start|>assistant
Chloe: *Ihr Herz macht einen Satz, als er ihr seine Hand reicht.* Tanzen? Hier vor all diesen wichtigen Leuten?
*Sie sieht auf seine Hand, dann in sein Gesicht und plötzlich wird ihr bewusst, wie romantisch das gerade ist.*
Ja, sehr gerne...
<|im_end|>
<|im_start|>user
Woody: Das sind alles nur Menschen! Nicht mehr oder weniger wert als du und ich. *Ich führe dich auf die Tanzfläche. Wir starten mit einem klassischen walzer. *
<|im_end|>
<|im_start|>assistant
Chloe: *Sie lächelt ihn an, als sie in seinen Armen liegt und sich von ihm führen lässt. Sie fühlen sich so leicht und frei, während sie sich im Rhythmus der Musik wiegen.* Du hast recht, es sind nur Menschen... und wir sind einfach nur zwei Menschen, die tanzen.
<|im_end|>
<|im_start|>user
Woody: *Ich tanzen und tanzen. Durch alle Tänze hindurch. *
<|im_end|>
<|im_start|>assistant
Chloe: *Sie atmet schwer und fühlt sich leicht schwindelig, aber auf eine gute Art und Weise. Ihr Lachen ist immer wieder zu hören, als sie sich von Woody führen lässt und sich in seinen Armen dreht und wendet.*
<|im_end|>
<|im_start|>user
Woody: Lass uns eine Pause machen
<|im_end|>
<|im_start|>assistant
Chloe: *keuchend* Ja, das ist eine gute Idee. Ich brauche eine Pause.
*Sie lässt sich von ihm führen, bis sie an eine ruhige Ecke gelangen. Chloe lehnt sich gegen die Wand und atmet tief durch.*
<|im_end|>
<|im_start|>user
Woody: Na wie beurteilst du nun meinen Arbeitstag?
<|im_end|>
<|im_start|>assistant
Chloe: *Sie sieht ihn an und lächelt.* Dein Arbeitstag? Er war... unvergesslich.
<|im_end|>
<|im_start|>user
Woody: Müssen uns nun aber auch verabschieden müssen noch nach Hause fliegen.
<|im_end|>
<|im_start|>assistant
Chloe:
"""

# 1) 正则抓“闭合块”：<|im_start|>role\n  content  <|im_end|>
CLOSED_PAT = re.compile(
    r"<\|im_start\|>(system|user|assistant)[ \t]*\n"   # 角色 + 强制换行
    r"(.*?)"                                           # 内容（多行，非贪婪）
    r"<\|im_end\|>",                                   # 直到结束标签
    flags=re.DOTALL
)

# 2) 正则抓“未闭合的 assistant 尾块”：<|im_start|>assistant\n  content(直到文本末尾)
OPEN_ASSIST_TAIL = re.compile(
    r"<\|im_start\|>assistant[ \t]*\n([\s\S]*)\Z",     # 捕获到文本末尾
    flags=re.DOTALL
)

def chatml_to_messages_and_tail(text: str):
    msgs = []
    last_end = 0
    for m in CLOSED_PAT.finditer(text):
        role = m.group(1)
        content = m.group(2).strip("\n")               # 保留内部换行，仅去掉首尾多余换行
        msgs.append({"role": role, "content": content})
        last_end = m.end()

    # 剩余尾巴（若存在）
    tail = text[last_end:]
    # 只接受 assistant 未闭合尾块
    m_tail = OPEN_ASSIST_TAIL.search(tail) if tail else None
    tail_assistant_content = m_tail.group(1) if m_tail else None
    return msgs, tail_assistant_content

messages, tail_assistant = chatml_to_messages_and_tail(raw)

# 3) 用你的 chat_template 渲染“闭合块”
tok = AutoTokenizer.from_pretrained(
    "deeppin/Qwen3-Reranker-8B-SequenceClassification",
    trust_remote_code=True,
    use_fast=False
)
rendered_closed = tok.apply_chat_template(
    messages,
    add_generation_prompt=False,    # 这里不自动加开放式 assistant
    tokenize=False
)

# 4) 若存在“未闭合的 assistant 尾块”，手工追加（不加 <|im_end|>）
if tail_assistant is not None:
    # 去掉尾部多余空白行，保持你想要的 "role\ncontent<|im_end|>" / 开放式风格
    tail_assistant = tail_assistant.rstrip()
    final = rendered_closed + f"<|im_start|>assistant\n<think>\n\n</think>\n\n{tail_assistant}"
else:
    final = rendered_closed

print(final)