95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
import os
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
|
|
from datasets import Dataset
|
|
from collections import defaultdict
|
|
|
|
# Konfiguracja
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
MODEL_NAME = "gpt2" # Tymczasowo używamy mniejszego modelu do testów
|
|
SPECIAL_TOKENS = ["[CITATION_START]", "[CITATION_END]"]
|
|
|
|
class SourceMapper:
|
|
def __init__(self):
|
|
self.source_to_idx = defaultdict(lambda: len(self.source_to_idx))
|
|
self.idx_to_source = {}
|
|
|
|
def add_source(self, source):
|
|
if source not in self.source_to_idx:
|
|
idx = self.source_to_idx[source]
|
|
self.idx_to_source[idx] = source
|
|
|
|
def prepare_simple_dataset():
|
|
# Przykładowe dane - zastąp rzeczywistymi danymi
|
|
return [
|
|
{
|
|
"text": "[CITATION_START] Kodeks Pracy, Art. 1 [CITATION_END] Tekst artykułu...",
|
|
"source_idx": 0
|
|
},
|
|
{
|
|
"text": "[CITATION_START] Kodeks Pracy, Art. 2 [CITATION_END] Inny tekst...",
|
|
"source_idx": 1
|
|
}
|
|
]
|
|
|
|
def main():
|
|
# Inicjalizacja
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
|
tokenizer.add_special_tokens({"additional_special_tokens": SPECIAL_TOKENS})
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
# Przygotowanie danych
|
|
source_mapper = SourceMapper()
|
|
data = prepare_simple_dataset()
|
|
|
|
# Tworzenie datasetu
|
|
dataset = Dataset.from_dict({
|
|
"text": [d["text"] for d in data],
|
|
"source_idx": [d["source_idx"] for d in data]
|
|
})
|
|
|
|
# Tokenizacja
|
|
def tokenize_function(examples):
|
|
tokenized = tokenizer(
|
|
examples["text"],
|
|
truncation=True,
|
|
padding="max_length",
|
|
max_length=128,
|
|
return_tensors="pt"
|
|
)
|
|
return {
|
|
"input_ids": tokenized["input_ids"].squeeze(),
|
|
"attention_mask": tokenized["attention_mask"].squeeze(),
|
|
"labels": tokenized["input_ids"].squeeze().clone(),
|
|
}
|
|
|
|
tokenized_dataset = dataset.map(tokenize_function, batched=True)
|
|
|
|
# Model
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
|
|
model.resize_token_embeddings(len(tokenizer))
|
|
|
|
# Konfiguracja treningu
|
|
training_args = TrainingArguments(
|
|
output_dir="./results",
|
|
num_train_epochs=1,
|
|
per_device_train_batch_size=2,
|
|
gradient_accumulation_steps=1,
|
|
learning_rate=2e-5,
|
|
logging_steps=1,
|
|
remove_unused_columns=False
|
|
)
|
|
|
|
# Trainer
|
|
trainer = Trainer(
|
|
model=model,
|
|
args=training_args,
|
|
train_dataset=tokenized_dataset,
|
|
)
|
|
|
|
# Rozpoczęcie treningu
|
|
print("Rozpoczęcie treningu...")
|
|
trainer.train()
|
|
|
|
if __name__ == "__main__":
|
|
main() |