dodano model herbert
This commit is contained in:
parent
d9541a9a28
commit
4f486f021b
|
|
@ -0,0 +1,119 @@
|
|||
import os
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
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, DataCollatorForLanguageModeling
|
||||
|
||||
# 1️⃣ Inicjalizacja modelu do embeddingów
|
||||
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
|
||||
|
||||
# 2️⃣ Dodanie dokumentów i embeddingów
|
||||
def read_documents_from_file(file_path):
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
content = file.read()
|
||||
articles = content.split('\n\n')
|
||||
documents = []
|
||||
for article in articles:
|
||||
if article.strip().startswith('Art.'):
|
||||
documents.append(article.strip())
|
||||
return documents
|
||||
#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?"
|
||||
#]
|
||||
file_path = './docs/kodekspracy.txt' # Zmień na właściwą ścieżkę
|
||||
documents = read_documents_from_file(file_path)
|
||||
embeddings = embed_model.encode(documents)
|
||||
|
||||
# 3️⃣ Inicjalizacja FAISS i dodanie wektorów
|
||||
dim = embeddings.shape[1]
|
||||
index = faiss.IndexFlatL2(dim)
|
||||
index.add(np.array(embeddings, dtype=np.float32))
|
||||
|
||||
# 4️⃣ Przygotowanie danych treningowych
|
||||
def create_training_data():
|
||||
data = {
|
||||
"text": documents,
|
||||
"embedding": embeddings.tolist()
|
||||
}
|
||||
return Dataset.from_dict(data)
|
||||
|
||||
dataset = create_training_data()
|
||||
|
||||
# Podział danych na treningowe i ewaluacyjne
|
||||
split_dataset = dataset.train_test_split(test_size=0.25)
|
||||
train_dataset = split_dataset["train"]
|
||||
eval_dataset = split_dataset["test"]
|
||||
|
||||
# 5️⃣ Ładowanie modelu Gemma 2B
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model_name = "pczarnik/herbert-base-ner"
|
||||
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
|
||||
# 6️⃣ Konfiguracja LoRA
|
||||
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
|
||||
max_length = 384
|
||||
|
||||
def tokenize_function(examples):
|
||||
return tokenizer(
|
||||
examples["text"],
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
max_length=max_length
|
||||
)
|
||||
|
||||
tokenized_train = train_dataset.map(tokenize_function, batched=True)
|
||||
tokenized_eval = eval_dataset.map(tokenize_function, batched=True)
|
||||
|
||||
# 8️⃣ Parametry treningu
|
||||
training_args = TrainingArguments(
|
||||
output_dir="./results",
|
||||
eval_strategy="steps", # Ewaluacja co określoną liczbę kroków
|
||||
eval_steps=500, # Ewaluacja co 500 kroków
|
||||
save_strategy="steps", # Zapis modelu co określoną liczbę kroków
|
||||
save_steps=500, # Zapis modelu co 500 kroków
|
||||
learning_rate=1e-5,
|
||||
per_device_train_batch_size=2,
|
||||
per_device_eval_batch_size=2,
|
||||
num_train_epochs=16,
|
||||
weight_decay=0.01,
|
||||
load_best_model_at_end=True, # Wczytaj najlepszy model na końcu
|
||||
metric_for_best_model="loss", # Kryterium wyboru najlepszego modelu
|
||||
greater_is_better=False, # Niższy loss = lepszy model
|
||||
)
|
||||
|
||||
# 9️⃣ Data Collator
|
||||
data_collator = DataCollatorForLanguageModeling(
|
||||
tokenizer=tokenizer,
|
||||
mlm=False
|
||||
)
|
||||
|
||||
# 🔟 Trening modelu
|
||||
trainer = Trainer(
|
||||
model=model,
|
||||
args=training_args,
|
||||
train_dataset=tokenized_train,
|
||||
eval_dataset=tokenized_eval, # Dodany zestaw ewaluacyjny
|
||||
data_collator=data_collator,
|
||||
)
|
||||
|
||||
trainer.train()
|
||||
|
||||
# 1️⃣1️⃣ Zapis modelu
|
||||
model.save_pretrained("./trained_model/gemma")
|
||||
tokenizer.save_pretrained("./trained_model/gemma")
|
||||
|
||||
print("✅ Model został wytrenowany i zapisany!")
|
||||
Loading…
Reference in New Issue