102 lines
2.8 KiB
Python
102 lines
2.8 KiB
Python
import os
|
||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||
|
||
import torch
|
||
import weaviate
|
||
from weaviate.classes.config import Property, DataType, Configure
|
||
from weaviate.classes.query import Query
|
||
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq
|
||
from datasets import Dataset
|
||
|
||
# 1️⃣ Połączenie z bazą Weaviate
|
||
client = weaviate.WeaviateClient(
|
||
connection_params=weaviate.ConnectionParams.from_params(
|
||
http_host="weaviate",
|
||
http_port=8080,
|
||
http_secure=False,
|
||
grpc_host="weaviate",
|
||
grpc_port=50051,
|
||
grpc_secure=False,
|
||
)
|
||
)
|
||
|
||
# 2️⃣ Pobranie dokumentów z bazy Weaviate
|
||
collection_name = "Documents"
|
||
query = Query(collection_name).limit(1000)
|
||
result = client.query.run(query)
|
||
|
||
documents = []
|
||
file_names = []
|
||
|
||
for item in result[collection_name]['objects']:
|
||
documents.append(item['properties']['content'])
|
||
file_names.append(item['properties']['fileName'])
|
||
|
||
# 3️⃣ Tworzenie datasetu
|
||
training_data = {
|
||
"text": documents,
|
||
"file_name": file_names
|
||
}
|
||
dataset = Dataset.from_dict(training_data)
|
||
|
||
# Podział na zestaw treningowy i ewaluacyjny
|
||
split_dataset = dataset.train_test_split(test_size=0.25)
|
||
train_dataset = split_dataset["train"]
|
||
eval_dataset = split_dataset["test"]
|
||
|
||
# 4️⃣ Ładowanie modelu Multislav
|
||
model_name = "allegro/multislav-5lang"
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
|
||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||
|
||
# 5️⃣ Tokenizacja
|
||
max_length = 512
|
||
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)
|
||
|
||
# 6️⃣ Parametry treningu
|
||
training_args = TrainingArguments(
|
||
output_dir="./results",
|
||
evaluation_strategy="steps",
|
||
eval_steps=500,
|
||
save_strategy="steps",
|
||
save_steps=500,
|
||
learning_rate=2e-5,
|
||
per_device_train_batch_size=4,
|
||
per_device_eval_batch_size=4,
|
||
num_train_epochs=10,
|
||
weight_decay=0.01,
|
||
load_best_model_at_end=True,
|
||
metric_for_best_model="loss",
|
||
greater_is_better=False,
|
||
)
|
||
|
||
# 7️⃣ Data Collator
|
||
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
|
||
|
||
# 8️⃣ Trening
|
||
trainer = Trainer(
|
||
model=model,
|
||
args=training_args,
|
||
train_dataset=tokenized_train,
|
||
eval_dataset=tokenized_eval,
|
||
data_collator=data_collator,
|
||
)
|
||
|
||
trainer.train()
|
||
|
||
# 9️⃣ Zapis modelu
|
||
model.save_pretrained("./trained_model/multislav")
|
||
tokenizer.save_pretrained("./trained_model/multislav")
|
||
|
||
print("✅ Model został wytrenowany i zapisany!")
|