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()