77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
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!") |