Finetuning.py

Finetuning.py

"""
Fine-tuning LoRA al unui model mic (Mistral 7B / Llama 3 8B) pe textul cărților.
Rezultat: un model care scrie și folosește jargonul de epocă din cărțile tale.

Instalare:
    pip install "transformers>=4.44" "trl>=0.9" "peft>=0.12" \
                "datasets>=2.20" "bitsandbytes>=0.43" accelerate sentencepiece
Rulare:
    python finetuning_carti.py
GPU recomandat: minim ~12 GB VRAM (cu quantizare 4-bit). Fără GPU nu e practic.
"""

import os
import glob
import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, prepare_model_for_kbit_training
from trl import SFTTrainer

# ----------------------------------------------------------------------
# 1. CONFIGURARE
# ----------------------------------------------------------------------
MODEL_NAME = "mistralai/Mistral-7B-v0.3"   # sau "meta-llama/Meta-Llama-3-8B"
TEXTS_DIR = "./carti_txt"                   # folder cu fișierele .txt ale cărților
OUTPUT_DIR = "./model_carti_lora"
MAX_SEQ_LEN = 2048                          # lungimea ferestrei de context la antrenare
BLOCK_OVERLAP = 128                         # suprapunere între blocuri (păstrează continuitatea)

# ----------------------------------------------------------------------
# 2. ÎNCĂRCAREA ȘI SEGMENTAREA TEXTELOR
#    Cărțile se taie în blocuri de dimensiune fixă, cu mică suprapunere.
# ----------------------------------------------------------------------
def incarca_texte(director):
    texte = []
    for cale in glob.glob(os.path.join(director, "*.txt")):
        with open(cale, "r", encoding="utf-8") as f:
            continut = f.read().strip()
        if continut:
            texte.append(continut)
    if not texte:
        raise SystemExit(f"Niciun .txt găsit în {director}")
    return texte


def segmenteaza(texte, tokenizer, lungime, suprapunere):
    blocuri = []
    for text in texte:
        ids = tokenizer(text, add_special_tokens=False)["input_ids"]
        pas = lungime - suprapunere
        for start in range(0, len(ids), pas):
            fragment = ids[start:start + lungime]
            if len(fragment) < 64:          # aruncă cozile prea scurte
                continue
            blocuri.append(tokenizer.decode(fragment))
    return Dataset.from_dict({"text": blocuri})


# ----------------------------------------------------------------------
# 3. MODEL + TOKENIZER (quantizat 4-bit pentru VRAM redus)
# ----------------------------------------------------------------------
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False

# ----------------------------------------------------------------------
# 4. CONFIGURARE LoRA (antrenăm doar adaptoare mici, nu tot modelul)
# ----------------------------------------------------------------------
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
)

# ----------------------------------------------------------------------
# 5. DATE
# ----------------------------------------------------------------------
texte = incarca_texte(TEXTS_DIR)
dataset = segmenteaza(texte, tokenizer, MAX_SEQ_LEN, BLOCK_OVERLAP)
print(f"Blocuri de antrenare: {len(dataset)}")

# ----------------------------------------------------------------------
# 6. ANTRENARE
# ----------------------------------------------------------------------
args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,          # batch efectiv = 16
    num_train_epochs=3,
    learning_rate=2e-4,
    bf16=True,
    logging_steps=10,
    save_strategy="epoch",
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    optim="paged_adamw_8bit",
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LEN,
    packing=True,                           # împachetează blocuri scurte → antrenare eficientă
)

trainer.train()
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
print(f"Adaptoarele LoRA au fost salvate în {OUTPUT_DIR}")

# ----------------------------------------------------------------------
# 7. TESTARE RAPIDĂ DUPĂ ANTRENARE
# ----------------------------------------------------------------------
model.config.use_cache = True
model.eval()
prompt = "Despre tratamentul frigurilor de baltă, autorii vechi recomandă"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    out = model.generate(**inputs, max_new_tokens=200, temperature=0.7, do_sample=True)
print(tokenizer.decode(out[0], skip_special_tokens=True))