import torch
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from datasets import Dataset
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer

# 1️⃣ Inicjalizacja modelu do embeddingów
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

# 2️⃣ Dodanie dokumentów i embeddingów
documents = [
    "Jak założyć firmę w Polsce?", 
    "Jak rozliczyć podatek VAT?", 
    "Procedura składania reklamacji w e-sklepie.",
    "Jakie dokumenty są potrzebne do rejestracji działalności?"
]
embeddings = embed_model.encode(documents)

# 3️⃣ Inicjalizacja FAISS i dodanie wektorów
dim = embeddings.shape[1]  # Wymiary wektorów
index = faiss.IndexFlatL2(dim)  # Tworzymy indeks FAISS dla metryki L2
index.add(np.array(embeddings, dtype=np.float32))  # Dodajemy wektory do indeksu FAISS

# 4️⃣ Przygotowanie danych treningowych
def create_training_data():
    # Pobranie dokumentów (możesz połączyć je z odpowiednimi embeddingami, jeśli trzeba)
    data = {
        "text": documents,
        "embedding": embeddings.tolist()
    }
    return Dataset.from_dict(data)

dataset = create_training_data()

# 5️⃣ Ładowanie modelu Gemma 2 7B
device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "google/gemma-7b"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 6️⃣ Konfiguracja LoRA dla efektywnego treningu
lora_config = LoraConfig(
    r=8, lora_alpha=32, lora_dropout=0.1, bias="none", task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

# 7️⃣ Tokenizacja danych
def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# 8️⃣ Parametry treningu
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=2,
    num_train_epochs=3,
    logging_dir="./logs",
    save_strategy="epoch"
)

# 9️⃣ Trening modelu
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
)

trainer.train()

# 🔟 Zapisanie dostrojonego modelu
model.save_pretrained("./trained_model/gemma")
tokenizer.save_pretrained("./trained_model/gemma")

print("✅ Model został wytrenowany i zapisany!")