File size: 3,833 Bytes
3eecd6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """
Chat Template Helper
Auto-detect format: native chat template (Qwen/Llama/Mistral) atau simple (GPT-2).
Dipakai oleh chat.py, rag_module.py, dan trainer.py.
"""
FORMAT_NATIVE = "native"
FORMAT_SIMPLE = "simple"
DEFAULT_SYSTEM_PROMPT = "Anda adalah asisten AI yang ramah dan informatif."
def detect_format(tokenizer) -> str:
"""Deteksi format yang didukung tokenizer."""
if getattr(tokenizer, "chat_template", None):
return FORMAT_NATIVE
return FORMAT_SIMPLE
def format_inference_prompt(tokenizer, user_message, system_message=None, history=None):
"""Format prompt untuk inference/chat."""
fmt = detect_format(tokenizer)
system_message = system_message or DEFAULT_SYSTEM_PROMPT
if fmt == FORMAT_NATIVE:
messages = [{"role": "system", "content": system_message}]
if history:
for turn in history[-5:]:
u = (turn.get("user") or "").strip()
a = (turn.get("ai") or "").strip()
if u and a:
messages.append({"role": "user", "content": u})
messages.append({"role": "assistant", "content": a})
messages.append({"role": "user", "content": user_message})
return tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
# Fallback: format lama "User: / AI:"
parts = [system_message]
if history:
for turn in history[-2:]:
u = (turn.get("user") or "").strip()
a = (turn.get("ai") or "").strip()
if u and a:
parts.append(f"User: {u}")
parts.append(f"AI: {a}")
parts.append(f"User: {user_message}")
parts.append("AI:")
return "\n".join(parts)
def format_training_sample(tokenizer, user_message, ai_response, system_message=None):
"""Format 1 Q&A untuk training."""
fmt = detect_format(tokenizer)
system_message = system_message or DEFAULT_SYSTEM_PROMPT
if fmt == FORMAT_NATIVE:
messages = [
{"role": "system", "content": system_message},
{"role": "user", "content": user_message},
{"role": "assistant", "content": ai_response},
]
return tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=False
)
eos = tokenizer.eos_token or ""
return f"User: {user_message}\nAI: {ai_response}{eos}"
def tokenize_for_training(tokenizer, text, max_length):
"""Tokenize + prompt masking untuk training."""
enc = tokenizer(
text,
truncation=True,
max_length=max_length,
padding=False,
return_tensors=None,
)
ids = enc["input_ids"]
labels = list(ids)
fmt = detect_format(tokenizer)
marker = "<|im_start|>assistant\n" if fmt == FORMAT_NATIVE else "AI:"
marker_ids = tokenizer.encode(marker, add_special_tokens=False)
mlen = len(marker_ids)
if mlen > 0 and len(ids) > mlen:
pos = -1
for i in range(len(ids) - mlen, -1, -1):
if ids[i:i + mlen] == marker_ids:
pos = i + mlen
break
if pos > 0:
for i in range(min(pos, len(labels))):
labels[i] = -100
return {
"input_ids": ids,
"attention_mask": enc["attention_mask"],
"labels": labels,
}
def tokenize_batch_for_training(examples, tokenizer, max_length):
"""Tokenize batch (untuk dataset.map)."""
result = {"input_ids": [], "attention_mask": [], "labels": []}
for text in examples["text"]:
out = tokenize_for_training(tokenizer, text, max_length)
result["input_ids"].append(out["input_ids"])
result["attention_mask"].append(out["attention_mask"])
result["labels"].append(out["labels"])
return result |