tmp allegro
This commit is contained in:
parent
cd535b4fe3
commit
029662e9d1
123
allegro.py
123
allegro.py
|
|
@ -1,122 +1,37 @@
|
||||||
import os
|
from transformers import MarianForCausalLM, MarianTokenizer, Trainer, TrainingArguments
|
||||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
from datasets import load_dataset
|
||||||
|
|
||||||
from weaviate.connect import ConnectionParams
|
# Załaduj model i tokenizer
|
||||||
import weaviate
|
|
||||||
|
|
||||||
import torch
|
|
||||||
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, MarianForCausalLM, MarianTokenizer
|
|
||||||
|
|
||||||
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
|
|
||||||
|
|
||||||
# 1️⃣ Połączenie z Weaviate
|
|
||||||
client = weaviate.WeaviateClient(
|
|
||||||
connection_params=ConnectionParams.from_params(
|
|
||||||
http_host="weaviate",
|
|
||||||
http_port=8080,
|
|
||||||
http_secure=False,
|
|
||||||
grpc_host="weaviate",
|
|
||||||
grpc_port=50051,
|
|
||||||
grpc_secure=False,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
client.connect()
|
|
||||||
|
|
||||||
collection = client.collections.get("Document")
|
|
||||||
|
|
||||||
# 2️⃣ Pobranie dokumentów z Weaviate
|
|
||||||
def fetch_documents():
|
|
||||||
response = collection.query.fetch_objects()
|
|
||||||
documents = []
|
|
||||||
for o in response.objects:
|
|
||||||
file_name = o.properties.get("fileName", "unknown_file")
|
|
||||||
content = o.properties.get("content", "")
|
|
||||||
if content:
|
|
||||||
documents.append(f"fileName: {file_name}, content: {content}")
|
|
||||||
print(f"fileName: {file_name}")
|
|
||||||
return documents
|
|
||||||
#return documents
|
|
||||||
|
|
||||||
documents = fetch_documents()
|
|
||||||
|
|
||||||
embeddings = embed_model.encode(documents)
|
|
||||||
|
|
||||||
dim = embeddings.shape[1]
|
|
||||||
#index = faiss.IndexFlatL2(dim)
|
|
||||||
#index.add(np.array(embeddings, dtype=np.float32))
|
|
||||||
|
|
||||||
def create_training_data():
|
|
||||||
data = {
|
|
||||||
"text": documents,
|
|
||||||
"embedding": embeddings.tolist()
|
|
||||||
}
|
|
||||||
return Dataset.from_dict(data)
|
|
||||||
|
|
||||||
dataset = create_training_data()
|
|
||||||
|
|
||||||
split_dataset = dataset.train_test_split(test_size=0.25)
|
|
||||||
train_dataset = split_dataset["train"]
|
|
||||||
eval_dataset = split_dataset["test"]
|
|
||||||
|
|
||||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
||||||
model_name = "allegro/multislav-5lang"
|
model_name = "allegro/multislav-5lang"
|
||||||
model = MarianForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to(device)
|
model = MarianForCausalLM.from_pretrained(model_name)
|
||||||
tokenizer = MarianForCausalLM.from_pretrained(model_name)
|
tokenizer = MarianTokenizer.from_pretrained(model_name)
|
||||||
|
|
||||||
lora_config = LoraConfig(
|
# Załaduj dane (przykład dla tłumaczenia z języka rumuńskiego na angielski)
|
||||||
r=8, lora_alpha=32, lora_dropout=0.1, bias="none", task_type="CAUSAL_LM"
|
dataset = load_dataset("wmt16", "ro-en")
|
||||||
)
|
|
||||||
model = get_peft_model(model, lora_config)
|
|
||||||
|
|
||||||
max_length = 384
|
|
||||||
|
|
||||||
|
# Przetwórz dane do formatu odpowiedniego dla modelu
|
||||||
def tokenize_function(examples):
|
def tokenize_function(examples):
|
||||||
return tokenizer(
|
return tokenizer(examples['translation'], truncation=True, padding='max_length', max_length=128)
|
||||||
examples["text"],
|
|
||||||
padding="max_length",
|
|
||||||
truncation=True,
|
|
||||||
max_length=max_length
|
|
||||||
)
|
|
||||||
|
|
||||||
tokenized_train = train_dataset.map(tokenize_function, batched=True)
|
tokenized_datasets = dataset.map(tokenize_function, batched=True)
|
||||||
tokenized_eval = eval_dataset.map(tokenize_function, batched=True)
|
|
||||||
|
|
||||||
|
# Skonfiguruj trenera
|
||||||
training_args = TrainingArguments(
|
training_args = TrainingArguments(
|
||||||
output_dir="./results",
|
output_dir="./results",
|
||||||
eval_strategy="steps", # Ewaluacja co określoną liczbę kroków
|
evaluation_strategy="epoch",
|
||||||
eval_steps=500, # Ewaluacja co 500 kroków
|
learning_rate=5e-5,
|
||||||
save_strategy="steps", # Zapis modelu co określoną liczbę kroków
|
per_device_train_batch_size=4,
|
||||||
save_steps=500, # Zapis modelu co 500 kroków
|
per_device_eval_batch_size=4,
|
||||||
learning_rate=1e-5,
|
num_train_epochs=3,
|
||||||
per_device_train_batch_size=2,
|
|
||||||
per_device_eval_batch_size=2,
|
|
||||||
num_train_epochs=16,
|
|
||||||
weight_decay=0.01,
|
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
|
|
||||||
)
|
|
||||||
|
|
||||||
data_collator = DataCollatorForLanguageModeling(
|
|
||||||
tokenizer=tokenizer,
|
|
||||||
mlm=False
|
|
||||||
)
|
)
|
||||||
|
|
||||||
trainer = Trainer(
|
trainer = Trainer(
|
||||||
model=model,
|
model=model,
|
||||||
args=training_args,
|
args=training_args,
|
||||||
train_dataset=tokenized_train,
|
train_dataset=tokenized_datasets["train"],
|
||||||
eval_dataset=tokenized_eval, # Dodany zestaw ewaluacyjny
|
eval_dataset=tokenized_datasets["test"],
|
||||||
data_collator=data_collator,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Trening modelu
|
||||||
trainer.train()
|
trainer.train()
|
||||||
|
|
||||||
model.save_pretrained("./trained_model/gemma")
|
|
||||||
tokenizer.save_pretrained("./trained_model/gemma")
|
|
||||||
|
|
||||||
print("✅ Model został wytrenowany i zapisany!")
|
|
||||||
Loading…
Reference in New Issue