Views
No views yet
1import torch
2from copy import deepcopy
3from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
4
5# ===== モデル =====
6model_path = "DataPilot/ArrowIdeative-13b-NeoBase-ZERO-llm-jp-v0.2"
7
8tokenizer = AutoTokenizer.from_pretrained(model_path)
9model = AutoModelForCausalLM.from_pretrained(
10 model_path,
11 device_map="auto",
12 torch_dtype=torch.bfloat16,
13)
14model.eval()
15
16system_prompt = """あなたは有能なアシスタントです。日本語で丁寧に答えてください。"""
17prompt = """CPUとGPUの違いについて教えてください。"""
18
19# (元コードのChatML形式を維持)
20text = f"""<|im_start|>system
21{system_prompt}<|im_end|>
22<|im_start|>user
23{prompt}<|im_end|>
24<|im_start|>assistant
25"""
26
27inputs = tokenizer(text, add_special_tokens=False, return_tensors="pt", return_token_type_ids=False).to(model.device)
28prompt_len = inputs["input_ids"].shape[1]
29
30# "<|im_end|>" のトークン列(1トークンとは限らないので列で扱う)
31stop_ids = tokenizer.encode("<|im_end|>", add_special_tokens=False)
32stop_ids = torch.tensor(stop_ids, device=model.device, dtype=inputs["input_ids"].dtype)
33
34class StopOnImEnd(StoppingCriteria):
35 def __init__(self, stop_ids_tensor: torch.Tensor):
36 super().__init__()
37 self.stop_ids = stop_ids_tensor
38
39 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
40 k = int(self.stop_ids.numel())
41 if k == 0 or input_ids.shape[1] < k:
42 return False
43 return torch.equal(input_ids[0, -k:], self.stop_ids)
44
45stopping_criteria = StoppingCriteriaList([StopOnImEnd(stop_ids)])
46
47# 既定EOSで止まらないようにする(= "<|im_end|>" のみで停止させる)
48gen_config = deepcopy(model.generation_config)
49gen_config.eos_token_id = None
50gen_config.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else model.config.eos_token_id
51
52with torch.inference_mode():
53 output = model.generate(
54 **inputs,
55 generation_config=gen_config,
56 stopping_criteria=stopping_criteria,
57 max_new_tokens=1024,
58 do_sample=True,
59 top_p=0.95,
60 temperature=0.5,
61 repetition_penalty=1.05,
62 )
63
64generated = tokenizer.decode(output[0, prompt_len:], skip_special_tokens=False)
65print(generated.split("<|im_end|>", 1)[0])
66<|im_end|>)の適切な出力とフォーマット準拠を評価None(マスク)として無視され学習に影響しない1@misc{arrowideative_13b_neobase_zero_llm_jp,
2 title = {ArrowIdeative-13b-NeoBase-ZERO-llm-jp},
3 author = {holy-fox},
4 year = {2026},
5}