76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
|
import torch
|
|||
|
|
import chromadb
|
|||
|
|
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 ChromaDB i modelu do embeddingów
|
|||
|
|
chroma_client = chromadb.PersistentClient(path="./chroma_db")
|
|||
|
|
collection = chroma_client.get_or_create_collection("my_embeddings")
|
|||
|
|
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).tolist()
|
|||
|
|
|
|||
|
|
for i, (doc, emb) in enumerate(zip(documents, embeddings)):
|
|||
|
|
collection.add(ids=[str(i)], documents=[doc], embeddings=[emb])
|
|||
|
|
|
|||
|
|
# 3️⃣ Przygotowanie danych treningowych
|
|||
|
|
def create_training_data():
|
|||
|
|
data = collection.get(include=["documents", "embeddings"])
|
|||
|
|
return Dataset.from_dict({
|
|||
|
|
"text": data["documents"],
|
|||
|
|
"embedding": data["embeddings"]
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
dataset = create_training_data()
|
|||
|
|
|
|||
|
|
# 4️⃣ Ł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)
|
|||
|
|
|
|||
|
|
# 5️⃣ 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)
|
|||
|
|
|
|||
|
|
# 6️⃣ Tokenizacja danych
|
|||
|
|
def tokenize_function(examples):
|
|||
|
|
return tokenizer(examples["text"], padding="max_length", truncation=True)
|
|||
|
|
|
|||
|
|
tokenized_dataset = dataset.map(tokenize_function, batched=True)
|
|||
|
|
|
|||
|
|
# 7️⃣ Parametry treningu
|
|||
|
|
training_args = TrainingArguments(
|
|||
|
|
output_dir="./results",
|
|||
|
|
per_device_train_batch_size=2,
|
|||
|
|
num_train_epochs=3,
|
|||
|
|
logging_dir="./logs",
|
|||
|
|
save_strategy="epoch"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 8️⃣ Trening modelu
|
|||
|
|
trainer = Trainer(
|
|||
|
|
model=model,
|
|||
|
|
args=training_args,
|
|||
|
|
train_dataset=tokenized_dataset,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
trainer.train()
|
|||
|
|
|
|||
|
|
# 9️⃣ Zapisanie dostrojonego modelu
|
|||
|
|
model.save_pretrained("./trained_model/gemma")
|
|||
|
|
tokenizer.save_pretrained("./trained_model/gemma")
|
|||
|
|
|
|||
|
|
print("✅ Model został wytrenowany i zapisany!")
|