init
This commit is contained in:
@@ -0,0 +1,519 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user