import os
import torch
from transformers import LlamaTokenizer, LlamaForCausalLM
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
from torch.utils.data import DataLoader
from tqdm import tqdm

# Настройки
MODEL_PATH = "model/"
LORA_RANK = 8
BATCH_SIZE = 1
EPOCHS = 3
LEARNING_RATE = 1e-4
MAX_LENGTH = 256
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Загрузка модели и токенизатора
model = LlamaForCausalLM.from_pretrained("model/",  gguf_file="Meta-Llama-3-8B-Instruct-Q4_K_M.gguf")
tokenizer = LlamaTokenizer.from_pretrained("model/", gguf_file="Meta-Llama-3-8B-Instruct-Q4_K_M.gguf")

# Настройка LoRA
lora_config = LoraConfig(
    r=LORA_RANK,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# Перемещение модели на GPU и включение смешанной точности
torch.cuda.empty_cache()
model.to(DEVICE)
model.enable_input_require_grads()
model = torch.compile(model)

# Загрузка данных
dataset = load_dataset("text", data_files="txt/*.txt")
train_dataloader = DataLoader(dataset["train"], batch_size=BATCH_SIZE, shuffle=True)

# Оптимизатор и планировщик скорости обучения
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_dataloader) * EPOCHS)

# Обучение
model.train()
for epoch in range(EPOCHS):
    for batch in tqdm(train_dataloader, desc=f"Epoch {epoch + 1}/{EPOCHS}"):
        inputs = tokenizer(batch["text"], padding=True, truncation=True, max_length=MAX_LENGTH, return_tensors="pt").to(DEVICE)
        outputs = model(**inputs, labels=inputs["input_ids"])
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

# Сохранение обученной модели
model.save_pretrained("lora/")

# Генерация стихотворения
prompt = "Напиши стихотворение про осень в Москве"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
output = model.generate(input_ids, max_length=200, num_return_sequences=1, temperature=0.7)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_text)    Converting and de-quantizing GGUF tensors...: 100%|██████████| 291/291 [01:27<00:00,  3.33it/s]
You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama.LlamaTokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message
Traceback (most recent call last):
  File "C:\Users\DM\PycharmProjects\pythonProject5\main.py", line 35, in <module>
    model.to(DEVICE)
  File "C:\Users\DM\PycharmProjects\pythonProject5\venv\lib\site-packages\torch\nn\modules\module.py", line 1174, in to
    return self._apply(convert)
  File "C:\Users\DM\PycharmProjects\pythonProject5\venv\lib\site-packages\torch\nn\modules\module.py", line 780, in _apply
    module._apply(fn)
  File "C:\Users\DM\PycharmProjects\pythonProject5\venv\lib\site-packages\torch\nn\modules\module.py", line 780, in _apply
    module._apply(fn)
  File "C:\Users\DM\PycharmProjects\pythonProject5\venv\lib\site-packages\torch\nn\modules\module.py", line 780, in _apply
    module._apply(fn)
  [Previous line repeated 4 more times]
  File "C:\Users\DM\PycharmProjects\pythonProject5\venv\lib\site-packages\torch\nn\modules\module.py", line 805, in _apply
    param_applied = fn(param)
  File "C:\Users\DM\PycharmProjects\pythonProject5\venv\lib\site-packages\torch\nn\modules\module.py", line 1160, in convert
    return t.to(
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 224.00 MiB. GPU 0 has a total capacity of 11.99 GiB of which 0 bytes is free. Of the allocated memory 11.27 GiB is allocated by PyTorch, and 794.00 KiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Process finished with exit code 1