modeling_youtu.py

-> 474 line

@check_model_inputs # <--- 这是导致问题的行
@auto_docstring
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
# ... other arguments ...
):

@check_model_inputs # <--- 已注释掉

@auto_docstring
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
attention_mask: Optional[torch.Tensor] = None,
# ... other arguments ...
):

import re
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 1. 配置模型
model_id = "./Youtu-LLM-2B"

# 2. 初始化Tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    trust_remote_code=True
)

# 3. 构建对话输入
prompt = "Hello"
messages = [{"role": "user", "content": prompt}]

# 使用 apply_chat_template 构建输入
input_ids = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
    enable_thinking=True
).to(model.device)

# 手动创建 attention_mask
attention_mask = torch.ones_like(input_ids).to(model.device)


# 4. 生成回复
# 在移除了错误的装饰器后，标准的调用方式应该可以正常工作
outputs = model.generate(
    input_ids=input_ids,
    attention_mask=attention_mask,
    max_new_tokens=512,
    do_sample=True,
    temperature=1.0,
    top_p=0.95,
    repetition_penalty=1.05,
    pad_token_id=tokenizer.eos_token_id # 添加此行以避免警告
)

# 5. 解析结果
full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)

def parse_reasoning(text):
    """提取<think>标签内的思考过程和之后的回答内容"""
    thought_pattern = r"<think>(.*?)</think>"
    match = re.search(thought_pattern, text, re.DOTALL)

    if match:
        thought = match.group(1).strip()
        answer = text.split("</think>")[-1].strip()
    else:
        thought = "(No explicit thought process generated)"
        answer = text
    return thought, answer

thought, final_answer = parse_reasoning(full_response)

print(f"\n{'='*20} Thought Process {'='*20}\n{thought}")
print(f"\n{'='*20} Final Answer {'='*20}\n{final_answer}")