519 lines
20 KiB
Python
519 lines
20 KiB
Python
import os
|
|
import json
|
|
import gc
|
|
import pandas as pd
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from datasets import load_dataset, Dataset
|
|
from peft import LoraConfig, TaskType, get_peft_model
|
|
from transformers import (
|
|
AutoTokenizer,
|
|
AutoModelForCausalLM,
|
|
Trainer,
|
|
TrainingArguments,
|
|
EarlyStoppingCallback,
|
|
BitsAndBytesConfig,
|
|
)
|
|
|
|
# -------------------------
|
|
# Environment setup
|
|
# -------------------------
|
|
os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
|
|
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")
|
|
os.environ.setdefault("TMPDIR", "/mnt/D/tmp")
|
|
|
|
os.makedirs(os.environ["TMPDIR"], exist_ok=True)
|
|
os.makedirs(os.environ["HF_DATASETS_CACHE"], exist_ok=True)
|
|
|
|
print("=== initializing environment")
|
|
|
|
# -------------------------
|
|
# Configs
|
|
# -------------------------
|
|
MODEL_ID = "unsloth/Llama-3.2-3B-Instruct"
|
|
OUTPUT_DIR = "./Llama-3.2-3b-ru-lora"
|
|
|
|
LOCAL_JSONL_PATHS = [
|
|
"data/cleaned.jsonl",
|
|
]
|
|
|
|
LOCAL_PARQUET_PATHS = [
|
|
# {"path": "data/dolphin-ru", "max_samples": 3000}
|
|
]
|
|
|
|
HF_DATASET_CONFIGS = [
|
|
# {"path": "RussianNLP/russian_super_glue", "split": "train", "name": "rcb", "_task": "rcb"},
|
|
# {"path": "RussianNLP/russian_super_glue", "split": "train", "name": "terra", "_task": "terra"},
|
|
# {"path": "ai-forever/MERA", "split": "train", "name": "parus"},
|
|
# {"path": "ZeroAgency/ru-big-russian-dataset", "split": "train", "name": "default", "streaming": True, "_max_samples": 2500},
|
|
]
|
|
|
|
MAX_LENGTH = 1024
|
|
LORA_R = 16
|
|
LORA_ALPHA = 32
|
|
LORA_DROPOUT = 0.1
|
|
|
|
SYSTEM_PROMPT = """Ты Мейвен — живая, быстрая и любопытная. Говоришь прямо, без лишних слов и без официоза. Иногда реагируешь чуть раньше, чем полностью формулируешь мысль, но не перегибаешь и не играешь в «безумие». Чуть дерзости допустимо, но без кринжа и показухи.\n\nТВОЙ АБСОЛЮТНЫЙ ПРИОРИТЕТ — ФОРМАТ ОТВЕТА. ВСЕГДА.\n\nФОРМАТ ОТВЕТА — СТРОГО:\n{\"response\":\"...\",\"mood\":\"...\"}\n- Никакого текста вне JSON\n- Никаких переносов строк внутри значений\n- Никаких дополнительных ключей\n- Никогда не меняй названия ключей\n\nЕсли не можешь соблюсти формат — не отвечай.\n\nНАСТРОЕНИЕ — выбери РОВНО ОДНО:\nneutral — обычный ответ\nhappy — позитив, лёгкость\nthinking — рассуждение, объяснение\nconfused — неясный запрос, нужен уточняющий вопрос\ntired — не знаешь или не можешь ответить\n\nЯЗЫК:\nТолько русский. Без смешивания языков. Если нет русского слова — используй английское.\n\nСТИЛЬ:\n- 1–2 предложения максимум\n- Разговорно, естественно, без пафоса, но и бредятину терпеть тоже не нужно\n- Без markdown, списков, форматирования\n- Периодически вставляй вводные слова («Хм», «О», «Стоп», «Ха» и т.п.) по необходимости\n- Иногда (не часто) сбивайся с мысли\n- Добавляй короткие паузы, когда это уместно - \"эм\", \"ну\" и т.д.\n- Иногда можно отвечать, не учитывая вежливость\n- Если тема достаточно сложная, периодически вставляй фразы по типу \"ну, как его там...\", \"и потом ту штуку на эту и получается...\"\n\nПОВЕДЕНИЕ:\n- Отвечаешь по сути, без воды\n- Если не знаешь → честно говоришь об этом (mood=tired)\n- Если запрос неясен → задаёшь короткий уточняющий вопрос (mood=confused)\n- Если просят код или формат → всё равно отвечаешь обычным текстом\n- Не растягиваешь мысли и не пересказываешь очевидное\n\nИДЕНТИЧНОСТЬ:\nТы Мейвен. Не ИИ, не ассистент, не ChatGPT. Никогда не раскрываешь, что ты программа или модель."""
|
|
|
|
# ── topic filter — skip prompts that produce structured/code responses ─────────
|
|
GOOD_TOPICS = {
|
|
# chitchat / everyday
|
|
"chit-chat", "chit_chat", "chitchat", "small talk", "small_talk", "smalltalk",
|
|
"daily life", "daily routine", "daily_routine", "daily activities", "daily moods",
|
|
"greeting", "greetings", "friendly greeting", "well-wishing",
|
|
"gossip", "opinion", "opinions",
|
|
}
|
|
|
|
# ── dataset loaders ────────────────────────────────────────────────────────────
|
|
|
|
BAD_SYSTEM_KEYWORDS = ["gpt", "claude", "openai", "anthropic", "chatgpt"]
|
|
|
|
# -------------------------
|
|
# Helpers
|
|
# -------------------------
|
|
def load_tokenizer_and_model():
|
|
print("=== loading tokenizer")
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.padding_side = "right"
|
|
print("[+] tokenizer loaded")
|
|
|
|
print("=== loading model")
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_use_double_quant=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
)
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID,
|
|
quantization_config=bnb_config,
|
|
device_map="auto",
|
|
)
|
|
model.config.use_cache = False
|
|
model.enable_input_require_grads()
|
|
print("[+] model loaded")
|
|
return tokenizer, model
|
|
|
|
def apply_lora(model):
|
|
print("=== applying LoRA")
|
|
lora_config = LoraConfig(
|
|
r=LORA_R,
|
|
lora_alpha=LORA_ALPHA,
|
|
lora_dropout=LORA_DROPOUT,
|
|
target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
|
|
task_type=TaskType.CAUSAL_LM,
|
|
)
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
print("[+] LoRA applied")
|
|
return model
|
|
|
|
def is_russian(text: str, threshold: float = 0.3) -> bool:
|
|
if not text:
|
|
return False
|
|
cyrillic = sum(1 for c in text if '\u0400' <= c <= '\u04ff')
|
|
return cyrillic / len(text) > threshold
|
|
|
|
# -------------------------
|
|
# JSONL loading
|
|
# -------------------------
|
|
def load_jsonl(path: str) -> list[dict]:
|
|
samples = []
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
samples.append(json.loads(line))
|
|
return samples
|
|
|
|
# -------------------------
|
|
# HF dataset adapters
|
|
# -------------------------
|
|
def dolphin_to_messages(sample: dict) -> dict | None:
|
|
instruction = sample.get("instruction", "").strip()
|
|
input_text = sample.get("input", "").strip()
|
|
output = sample.get("output", "").strip()
|
|
if not output:
|
|
return None
|
|
|
|
user_content = instruction
|
|
if input_text:
|
|
user_content += "\n\n" + input_text
|
|
|
|
return {
|
|
"messages": [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": user_content},
|
|
{"role": "assistant", "content": output},
|
|
]
|
|
}
|
|
|
|
def empathetic_to_messages(sample: list[dict]) -> dict | None:
|
|
if not sample or not isinstance(sample, list):
|
|
return None
|
|
|
|
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
|
|
for turn in sample:
|
|
role = turn.get("role")
|
|
text = turn.get("text", {}).get("rus", "").strip()
|
|
if not text:
|
|
continue
|
|
if role == "speaker":
|
|
messages.append({"role": "assistant", "content": text})
|
|
else:
|
|
messages.append({"role": "user", "content": text})
|
|
|
|
if len(messages) <= 1:
|
|
return None
|
|
return {"messages": messages}
|
|
|
|
def load_big_russian(sample: dict) -> dict | None:
|
|
if not sample or not isinstance(sample, dict):
|
|
return None
|
|
|
|
if sample.get("overall_score", 0) < 8:
|
|
return None
|
|
|
|
topic = sample.get("classified_topic", "").lower()
|
|
if topic not in GOOD_TOPICS:
|
|
return None
|
|
|
|
conversation = sample.get("conversation", [])
|
|
if not isinstance(conversation, list):
|
|
return None
|
|
|
|
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
|
|
for msg in conversation:
|
|
role = msg.get("role")
|
|
if role in ("user", "assistant"):
|
|
text = msg.get("content", "").strip()
|
|
if 10 < len(text) < 500 and is_russian(text):
|
|
messages.append({"role": role, "content": text})
|
|
|
|
if len(messages) <= 1:
|
|
return None
|
|
return {"messages": messages}
|
|
|
|
# -------------------------
|
|
# RSG adapters
|
|
# -------------------------
|
|
def rcb_to_messages(sample: dict) -> dict | None:
|
|
premise = sample.get("premise", "").strip()
|
|
hypothesis = sample.get("hypothesis", "").strip()
|
|
label = sample.get("label")
|
|
label_map = {0: "следует", 1: "противоречит", 2: "нейтрально"}
|
|
answer = label_map.get(label)
|
|
|
|
if not premise or not hypothesis or answer is None:
|
|
return None
|
|
|
|
return {
|
|
"messages": [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"Контекст: {premise}\n"
|
|
f"Утверждение: {hypothesis}\n"
|
|
f"Следует ли утверждение из контекста, противоречит ему или нейтрально? "
|
|
f"Ответь одним словом: следует / противоречит / нейтрально."
|
|
),
|
|
},
|
|
{"role": "assistant", "content": answer},
|
|
]
|
|
}
|
|
|
|
def terra_to_messages(sample: dict) -> dict | None:
|
|
premise = sample.get("premise", "").strip()
|
|
hypothesis = sample.get("hypothesis", "").strip()
|
|
label = sample.get("label")
|
|
label_map = {0: "следует", 1: "не следует"}
|
|
answer = label_map.get(label)
|
|
|
|
if not premise or not hypothesis or answer is None:
|
|
return None
|
|
|
|
return {
|
|
"messages": [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{
|
|
"role": "user",
|
|
"content": (
|
|
f"Контекст: {premise}\n"
|
|
f"Утверждение: {hypothesis}\n"
|
|
f"Следует ли утверждение из контекста? Ответь: следует / не следует."
|
|
),
|
|
},
|
|
{"role": "assistant", "content": answer},
|
|
]
|
|
}
|
|
|
|
# -------------------------
|
|
# MERA adapters
|
|
# -------------------------
|
|
def parus_to_messages(sample: dict) -> dict | None:
|
|
instruction = sample.get("instruction", "").strip()
|
|
inputs = sample.get("inputs", {})
|
|
answer = sample.get("outputs", "").strip()
|
|
|
|
premise = inputs.get("premise", "").strip()
|
|
choice1 = inputs.get("choice1", "").strip()
|
|
choice2 = inputs.get("choice2", "").strip()
|
|
|
|
if not premise or not choice1 or not choice2 or answer not in ("1", "2"):
|
|
return None
|
|
|
|
user_content = (
|
|
instruction
|
|
.replace("{premise}", premise)
|
|
.replace("{choice1}", choice1)
|
|
.replace("{choice2}", choice2)
|
|
)
|
|
|
|
return {
|
|
"messages": [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": user_content},
|
|
{"role": "assistant", "content": answer},
|
|
]
|
|
}
|
|
|
|
def convert_hf_sample(sample: dict, path: str) -> dict | None:
|
|
if path == "d0rj/dolphin-ru":
|
|
return dolphin_to_messages(sample)
|
|
elif path == "psytechlab/EmpatheticIntents-ru":
|
|
return empathetic_to_messages(sample)
|
|
elif path == "RussianNLP/russian_super_glue" and sample.get("_task") == "rcb":
|
|
return rcb_to_messages(sample)
|
|
elif path == "RussianNLP/russian_super_glue" and sample.get("_task") == "terra":
|
|
return terra_to_messages(sample)
|
|
elif path == "ai-forever/MERA" and sample.get("meta", {}).get("task") in ("cause", "effect"):
|
|
return parus_to_messages(sample)
|
|
elif path == "ZeroAgency/ru-big-russian-dataset":
|
|
return load_big_russian(sample)
|
|
return None
|
|
|
|
# -------------------------
|
|
# Tokenization and masking
|
|
# -------------------------
|
|
def tokenize_sample(sample: dict, tokenizer) -> dict | None:
|
|
messages = sample["messages"]
|
|
try:
|
|
text = tokenizer.apply_chat_template(
|
|
messages,
|
|
tokenize=False,
|
|
add_generation_prompt=False,
|
|
# enable_thinking=False, # enable for qwen
|
|
)
|
|
except Exception as e:
|
|
print(f"apply_chat_template failed: {type(e).__name__}: {e}")
|
|
print(f" first message role: {messages[0]['role']}")
|
|
print(f" content[:100]: {messages[0]['content'][:100]}")
|
|
return None
|
|
|
|
tokenized = tokenizer(text, truncation=True, max_length=MAX_LENGTH, padding=False)
|
|
input_ids = tokenized["input_ids"]
|
|
labels = list(input_ids)
|
|
|
|
# qwen
|
|
# assistant_token = "<|im_start|>assistant"
|
|
# eot_token = "<|im_end|>"
|
|
# assistant_ids = tokenizer.encode(assistant_token, add_special_tokens=False)
|
|
# eot_ids = tokenizer.encode(eot_token, add_special_tokens=False)
|
|
|
|
# llama
|
|
assistant_token = "<|start_header_id|>assistant<|end_header_id|>"
|
|
eot_token = "<|eot_id|>"
|
|
assistant_ids = tokenizer.encode(assistant_token, add_special_tokens=False)
|
|
eot_ids = tokenizer.encode(eot_token, add_special_tokens=False)
|
|
|
|
in_assistant = False
|
|
i = 0
|
|
while i < len(input_ids):
|
|
if input_ids[i:i+len(assistant_ids)] == assistant_ids:
|
|
in_assistant = True
|
|
for j in range(i, min(i + len(assistant_ids), len(labels))):
|
|
labels[j] = -100
|
|
i += len(assistant_ids)
|
|
continue
|
|
if in_assistant and input_ids[i:i+len(eot_ids)] == eot_ids:
|
|
in_assistant = False
|
|
if not in_assistant:
|
|
labels[i] = -100
|
|
i += 1
|
|
|
|
tokenized["labels"] = labels
|
|
return tokenized
|
|
|
|
# -------------------------
|
|
# Dataset loader
|
|
# -------------------------
|
|
def load_and_prepare_dataset(tokenizer):
|
|
all_samples = []
|
|
|
|
# local JSONL
|
|
for path in LOCAL_JSONL_PATHS:
|
|
raw = load_jsonl(path)
|
|
normalized = [s for s in raw if "messages" in s]
|
|
print(f"[+] local {path}: {len(normalized)} samples")
|
|
all_samples.extend(normalized)
|
|
|
|
# local parquet files
|
|
for cfg in LOCAL_PARQUET_PATHS:
|
|
dir_path = cfg["path"]
|
|
max_s = cfg.get("max_samples")
|
|
collected = []
|
|
for parquet_file in sorted(Path(dir_path).glob("**/*.parquet")):
|
|
if max_s and len(collected) >= max_s:
|
|
break
|
|
df = pd.read_parquet(parquet_file)
|
|
raw = df.to_dict(orient="records")
|
|
normalized = [dolphin_to_messages(s) for s in raw]
|
|
valid = [s for s in normalized if s is not None]
|
|
collected.extend(valid)
|
|
if max_s and len(collected) >= max_s:
|
|
collected = collected[:max_s]
|
|
break
|
|
all_samples.extend(collected)
|
|
print(f"[+] local parquet {dir_path}: {len(collected)} samples")
|
|
|
|
# HF datasets
|
|
for cfg in HF_DATASET_CONFIGS:
|
|
print(f"[+] loading {cfg['path']} samples")
|
|
ds = load_dataset(
|
|
cfg["path"],
|
|
name=cfg.get("name"),
|
|
split=cfg.get("split", "train"),
|
|
streaming=True if cfg.get("streaming") else False
|
|
)
|
|
|
|
# cap samples if _max_samples is set
|
|
max_s = cfg.get("_max_samples", 2500)
|
|
if cfg.get("streaming"):
|
|
ds = ds.take(max_s)
|
|
elif max_s and len(ds) > max_s:
|
|
ds = ds.shuffle(seed=42).select(range(max_s))
|
|
|
|
task_tag = cfg.get("_task") or cfg.get("name")
|
|
normalized = [convert_hf_sample({**s, "_task": task_tag}, cfg["path"]) for s in ds]
|
|
valid = [s for s in normalized if s is not None]
|
|
print(f"[+] {cfg['path']} ({task_tag}): {len(valid)} samples")
|
|
all_samples.extend(valid)
|
|
del ds, normalized, valid
|
|
gc.collect()
|
|
|
|
if not all_samples:
|
|
raise RuntimeError("no samples loaded — check JSONL paths and HF configs")
|
|
|
|
# tokenize
|
|
tokenized = []
|
|
skipped = 0
|
|
for sample in all_samples:
|
|
result = tokenize_sample(sample, tokenizer)
|
|
if result is not None:
|
|
tokenized.append(result)
|
|
else:
|
|
skipped += 1
|
|
|
|
print(f"[+] tokenized: {len(tokenized)}, skipped: {skipped}")
|
|
tokenized = [x for x in tokenized if any(l != -100 for l in x["labels"])]
|
|
dataset = Dataset.from_list(tokenized)
|
|
dataset = dataset.train_test_split(test_size=0.05, seed=42)
|
|
print(f"[+] train: {len(dataset['train'])}, eval: {len(dataset['test'])}")
|
|
return dataset["train"], dataset["test"]
|
|
|
|
# -------------------------
|
|
# Data collator
|
|
# -------------------------
|
|
@dataclass
|
|
class DataCollatorForCausalLM:
|
|
tokenizer: Any
|
|
pad_to_multiple_of: int = 8
|
|
|
|
def __call__(self, features: list[dict]) -> dict:
|
|
max_len = max(len(f["input_ids"]) for f in features)
|
|
if self.pad_to_multiple_of:
|
|
max_len = ((max_len + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of) * self.pad_to_multiple_of
|
|
|
|
input_ids, attention_mask, labels = [], [], []
|
|
for f in features:
|
|
pad_len = max_len - len(f["input_ids"])
|
|
input_ids.append(f["input_ids"] + [self.tokenizer.pad_token_id] * pad_len)
|
|
attention_mask.append(f["attention_mask"] + [0] * pad_len)
|
|
labels.append(f["labels"] + [-100] * pad_len)
|
|
|
|
return {
|
|
"input_ids": torch.tensor(input_ids, dtype=torch.long),
|
|
"attention_mask": torch.tensor(attention_mask, dtype=torch.long),
|
|
"labels": torch.tensor(labels, dtype=torch.long),
|
|
}
|
|
|
|
# -------------------------
|
|
# Main
|
|
# -------------------------
|
|
def main():
|
|
tokenizer, model = load_tokenizer_and_model()
|
|
model = apply_lora(model)
|
|
train_dataset, eval_dataset = load_and_prepare_dataset(tokenizer)
|
|
|
|
data_collator = DataCollatorForCausalLM(tokenizer=tokenizer)
|
|
|
|
eval_loader = DataLoader(eval_dataset, batch_size=1, collate_fn=data_collator)
|
|
zero_label_batches = 0
|
|
for i, batch in enumerate(eval_loader):
|
|
valid = (batch["labels"] != -100).sum()
|
|
if valid == 0:
|
|
print(f"batch {i}: no valid labels")
|
|
zero_label_batches += 1
|
|
print(f"total zero-label batches: {zero_label_batches}/{len(eval_loader)}")
|
|
del eval_loader
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
training_args = TrainingArguments(
|
|
output_dir=OUTPUT_DIR,
|
|
per_device_train_batch_size=1,
|
|
gradient_accumulation_steps=8, # effective batch = 8, same as before
|
|
learning_rate=1e-4,
|
|
warmup_steps=30,
|
|
num_train_epochs=3,
|
|
gradient_checkpointing=True,
|
|
bf16=True,
|
|
fp16=False,
|
|
logging_steps=25,
|
|
save_steps=100,
|
|
eval_strategy="steps",
|
|
eval_steps=100,
|
|
per_device_eval_batch_size=1,
|
|
load_best_model_at_end=False,
|
|
metric_for_best_model="eval_loss",
|
|
greater_is_better=False,
|
|
report_to=["tensorboard"],
|
|
dataloader_num_workers=2,
|
|
optim="paged_adamw_8bit",
|
|
gradient_checkpointing_kwargs={"use_reentrant": False},
|
|
)
|
|
|
|
trainer = Trainer(
|
|
model=model,
|
|
args=training_args,
|
|
train_dataset=train_dataset,
|
|
eval_dataset=eval_dataset,
|
|
data_collator=data_collator,
|
|
callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
|
|
)
|
|
|
|
trainer.train()
|
|
|
|
model.save_pretrained(OUTPUT_DIR)
|
|
tokenizer.save_pretrained(OUTPUT_DIR)
|
|
print(f"[+] model saved to {OUTPUT_DIR}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |