diff --git a/herbert.py b/herbert.py new file mode 100644 index 0000000..d4638e2 --- /dev/null +++ b/herbert.py @@ -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!") \ No newline at end of file