🇹🇭 OpenThai Chinda 4B

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "iapp/chinda-qwen3-4b"

# 加载分词器和模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)

# 准备模型输入
prompt = "อธิบายเกี่ยวกับปัญญาประดิษฐ์ให้ฟังหน่อย"
messages = [
    {"role": "user", "content": prompt}
]

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True  # 启用思考模式以获得更好的推理能力
)

model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 生成响应
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=1024,
    temperature=0.6,
    top_p=0.95,
    top_k=20,
    do_sample=True
)

output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# 解析思考内容（如果已启用）
try:
    # 查找 </think> 标记 (151668)
    index = len(output_ids) - output_ids[::-1].index(151668)
except ValueError:
    index = 0

thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

print("🧠 思考:", thinking_content)
print("💬 回复:", content)

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True  # 启用详细推理
)

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False  # 快速响应模式
)

# 基本模板结构
messages = [
    {"role": "system", "content": "You are a helpful Thai AI assistant."},
    {"role": "user", "content": "สวัสดีครับ"},
    {"role": "assistant", "content": "สวัสดีค่ะ! มีอะไรให้ช่วยเหลือบ้างคะ"},
    {"role": "user", "content": "ช่วยอธิบายเรื่อง AI ให้ฟังหน่อย"}
]

# 应用带思考模式的模板
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True
)

# 带思考控制的多轮对话
def create_conversation(messages, enable_thinking=True):
    # 如果不存在，添加系统消息
    if not messages or messages[0]["role"] != "system":
        system_msg = {
            "role": "system",
            "content": "คุณเป็น AI ผู้ช่วยที่ฉลาดและเป็นประโยชน์ พูดภาษาไทยได้อย่างเป็นธรรมชาติ"
        }
        messages = [system_msg] + messages

    # 应用聊天模板
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_thinking
    )

    return text

# 用法示例
conversation = [
    {"role": "user", "content": "คำนวณ 15 × 23 = ?"},
]

prompt = create_conversation(conversation, enable_thinking=True)

# 为复杂问题启用思考
messages = [
    {"role": "user", "content": "/think แก้สมการ: x² + 5x - 14 = 0"}
]

# 为快速响应禁用思考
messages = [
    {"role": "user", "content": "/no_think สวัสดี"}
]

def manage_context(messages, max_tokens=30000):
    """简单的上下文管理函数"""
    total_tokens = sum(len(tokenizer.encode(msg["content"])) for msg in messages)

    while total_tokens > max_tokens and len(messages) > 2:
        # 保留系统消息并删除最旧的用户/助手对
        if messages[1]["role"] == "user":
            messages.pop(1)  # 删除用户消息
            if len(messages) > 1 and messages[1]["role"] == "assistant":
                messages.pop(1)  # 删除对应的助手消息

        total_tokens = sum(len(tokenizer.encode(msg["content"])) for msg in messages)

    return messages