diff --git a/README.md b/README.md index 49cc8ef..19396be 100644 Binary files a/README.md and b/README.md differ diff --git a/configs/training_args_unsloth_.yaml b/configs/training_args_unsloth_.yaml new file mode 100644 index 0000000..f9d3a0c --- /dev/null +++ b/configs/training_args_unsloth_.yaml @@ -0,0 +1,89 @@ +top.booster: unsloth +top.checkpoint_path: [] +top.finetuning_type: lora +top.model_name: Custom +top.quantization_bit: '4' +top.quantization_method: bnb +top.rope_scaling: none +top.template: alpaca +train.additional_target: '' +train.apollo_rank: 16 +train.apollo_scale: 32 +train.apollo_target: all +train.apollo_update_interval: 200 +train.badam_mode: layer +train.badam_switch_interval: 50 +train.badam_switch_mode: ascending +train.badam_update_ratio: 0.05 +train.batch_size: 1 +train.compute_type: fp16 +train.create_new_adapter: false +train.cutoff_len: 1024 +train.dataset: +- alpaca_slovak_cleaned +train.dataset_dir: data +train.ds_offload: false +train.ds_stage: none +train.enable_thinking: false +train.extra_args: '{"optim": "adamw_8bit", "eval_steps": 1000, "eval_strategy": "steps", + "save_total_limit": 2}' +train.freeze_extra_modules: '' +train.freeze_language_model: false +train.freeze_multi_modal_projector: true +train.freeze_trainable_layers: 2 +train.freeze_trainable_modules: all +train.freeze_vision_tower: true +train.galore_rank: 16 +train.galore_scale: 2 +train.galore_target: all +train.galore_update_interval: 200 +train.gradient_accumulation_steps: 8 +train.hub_private_repo: false +train.image_max_pixels: 768*768 +train.image_min_pixels: 32*32 +train.learning_rate: 2e-4 +train.logging_steps: 5 +train.lora_alpha: 32 +train.lora_dropout: 0.05 +train.lora_rank: 16 +train.lora_target: q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj +train.loraplus_lr_ratio: 0 +train.lr_scheduler_type: cosine +train.mask_history: false +train.max_grad_norm: '1.0' +train.max_samples: '50000' +train.neat_packing: false +train.neftune_alpha: 0 +train.num_train_epochs: '1.0' +train.packing: false +train.ppo_score_norm: false +train.ppo_whiten_rewards: false +train.pref_beta: 0.1 +train.pref_ftx: 0 +train.pref_loss: sigmoid +train.project: huggingface +train.report_to: none +train.resize_vocab: false +train.reward_model: [] +train.save_steps: 1000 +train.swanlab_api_key: '' +train.swanlab_link: null +train.swanlab_mode: cloud +train.swanlab_project: llamafactory +train.swanlab_run_name: '' +train.swanlab_workspace: '' +train.trackio_space_id: trackio +train.train_on_prompt: false +train.training_stage: Supervised Fine-Tuning +train.use_apollo: false +train.use_badam: false +train.use_dora: false +train.use_galore: false +train.use_llama_pro: false +train.use_pissa: false +train.use_rslora: false +train.use_swanlab: false +train.val_size: 0.025 +train.video_max_pixels: 256*256 +train.video_min_pixels: 16*16 +train.warmup_steps: 150 \ No newline at end of file diff --git a/evaluation/lmeval_custom_tasks/sklegal.yaml b/evaluation/lmeval_custom_tasks/sklegal.yaml new file mode 100644 index 0000000..6dd2963 --- /dev/null +++ b/evaluation/lmeval_custom_tasks/sklegal.yaml @@ -0,0 +1,26 @@ +task: sklegal +repeats: 1 +metadata: + version: 1 +test_split: test +description: '' +doc_to_text: '{{prompt}} ' +metric_list: +- metric: acc + aggregation: mean + higher_is_better: true +- metric: acc_norm + aggregation: mean + higher_is_better: true +num_fewshot: 0 +output_type: multiple_choice +unsafe_code: false +dataset_path: TUKE-KEMT/slovak-legal-exam +doc_to_choice: '{{answers}}' +doc_to_target: label +target_delimiter: ' ' +fewshot_delimiter: ' + + + ' +should_decontaminate: false diff --git a/evaluation/lmeval_custom_tasks/skquad.yaml b/evaluation/lmeval_custom_tasks/skquad.yaml new file mode 100644 index 0000000..2546a99 --- /dev/null +++ b/evaluation/lmeval_custom_tasks/skquad.yaml @@ -0,0 +1,2 @@ +task: skquad +class: !function task.SQuAD2 diff --git a/evaluation/lmeval_custom_tasks/task.py b/evaluation/lmeval_custom_tasks/task.py new file mode 100644 index 0000000..c9e55f7 --- /dev/null +++ b/evaluation/lmeval_custom_tasks/task.py @@ -0,0 +1,245 @@ +""" +Know What You Don’t Know: Unanswerable Questions for SQuAD +https://arxiv.org/pdf/1806.03822.pdf + +Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, +consisting of questions posed by crowdworkers on a set of Wikipedia articles, +where the answer to every question is a segment of text, or span, from the +corresponding reading passage, or the question might be unanswerable. +SQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable +questions written adversarially by crowdworkers to look similar to answerable ones. +To do well on SQuAD2.0, systems must not only answer questions when possible, but +also determine when no answer is supported by the paragraph and abstain from answering. + +Homepage: https://rajpurkar.github.io/SQuAD-explorer/ +""" + +from functools import partial +from math import exp + +import datasets +from packaging import version + +from lm_eval.api.instance import Instance +from lm_eval.api.task import ConfigurableTask + + +_CITATION = """ +@misc{rajpurkar2018know, + title={Know What You Don't Know: Unanswerable Questions for SQuAD}, + author={Pranav Rajpurkar and Robin Jia and Percy Liang}, + year={2018}, + eprint={1806.03822}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} +""" + + +def _squad_metric(predictions, references): + import evaluate + + squad_metric = evaluate.load("squad_v2") + return squad_metric.compute(predictions=predictions, references=references) + + +def _squad_agg(key, items): + predictions, references = zip(*items) + + return _squad_metric(predictions=predictions, references=references).get(key, 0) + + +class SQuAD2(ConfigurableTask): + VERSION = 3 + DATASET_PATH = "TUKE-DeutscheTelekom/skquad" + DATASET_NAME = None + + def __init__(self, config=None): + super().__init__(config={"metadata": {"version": self.VERSION}}) + + # HF changed squad on us so we have to make sure we aren't running the old one + assert version.parse(datasets.__version__) >= version.parse("1.11.0"), ( + "datasets v1.11.0 or later required for SQuAD" + ) + + def has_training_docs(self): + return True + + def has_validation_docs(self): + return True + + def has_test_docs(self): + return False + + def training_docs(self): + return self.dataset["train"] + + def validation_docs(self): + return self.dataset["validation"] + + def doc_to_text(self, doc): + return ( + "Title: " + + doc["title"] + + "\n\n" + + "Background: " + + doc["context"] + + "\n\n" + + "Question: " + + doc["question"] + + "\n\n" + + "Answer:" + ) + + def should_decontaminate(self): + return True + + def doc_to_decontamination_query(self, doc): + return doc["context"] + + def doc_to_target(self, doc): + answer_list = doc["answers"]["text"] + if len(answer_list) > 0: + answer = answer_list[0] + else: + answer = "unanswerable" + return " " + answer + + def construct_requests( + self, doc, ctx, chat_template=None, apply_chat_template=False, **kwargs + ): + """Uses RequestFactory to construct Requests and returns an iterable of + Requests which will be sent to the LM. + + :param doc: + The document as returned from training_docs, validation_docs, or test_docs. + :param ctx: str + The context string, generated by fewshot_context. This includes the natural + language description, as well as the few shot examples, and the question + part of the document for `doc`. + """ + + return [ + Instance( + request_type="generate_until", + doc=doc, + arguments=(ctx, {"until": ["\n"]}), + idx=0, + **kwargs, + ), + Instance( + request_type="loglikelihood", + doc=doc, + arguments=(ctx, " " + "unanswerable"), + idx=0, + **kwargs, + ), + ] + + def process_results(self, doc, results): + """Take a single document and the LM results and evaluates, returning a + dict where keys are the names of submetrics and values are the values of + the metric for that one document + + :param doc: + The document as returned from training_docs, validation_docs, or test_docs. + :param results: + The results of the requests created in construct_requests. + """ + + continuation, (logprob_unanswerable, _) = results + + no_answer_probability = exp(logprob_unanswerable) + + predictions = { + "id": doc["id"], + "prediction_text": continuation, + "no_answer_probability": no_answer_probability, + } + + references = { + "id": doc["id"], + "answers": doc["answers"], + } + + return { + "exact": ( + predictions, + references, + ), # Exact match (the normalized answer exactly match the gold answer) + "f1": ( + predictions, + references, + ), # The F-score of predicted tokens versus the gold answer + "HasAns_exact": ( + predictions, + references, + ), # Exact match (the normalized answer exactly match the gold answer) + "HasAns_f1": ( + predictions, + references, + ), # The F-score of predicted tokens versus the gold answer + "NoAns_exact": ( + predictions, + references, + ), # Exact match (the normalized answer exactly match the gold answer) + "NoAns_f1": ( + predictions, + references, + ), # The F-score of predicted tokens versus the gold answer + "best_exact": ( + predictions, + references, + ), # Best exact match (with varying threshold) + "best_f1": (predictions, references), # Best F1 (with varying threshold) + } + + def aggregation(self): + """ + :returns: {str: [float] -> float} + A dictionary where keys are the names of submetrics and values are + functions that aggregate a list of metrics + """ + return { + "exact": partial( + _squad_agg, "exact" + ), # Exact match (the normalized answer exactly match the gold answer) + "f1": partial( + _squad_agg, "f1" + ), # The F-score of predicted tokens versus the gold answer + "HasAns_exact": partial( + _squad_agg, "HasAns_exact" + ), # Exact match (the normalized answer exactly match the gold answer) + "HasAns_f1": partial( + _squad_agg, "HasAns_f1" + ), # The F-score of predicted tokens versus the gold answer + "NoAns_exact": partial( + _squad_agg, "NoAns_exact" + ), # Exact match (the normalized answer exactly match the gold answer) + "NoAns_f1": partial( + _squad_agg, "NoAns_f1" + ), # The F-score of predicted tokens versus the gold answer + "best_exact": partial( + _squad_agg, "best_exact" + ), # Best exact match (with varying threshold) + "best_f1": partial( + _squad_agg, "best_f1" + ), # Best F1 (with varying threshold) + } + + def higher_is_better(self): + """ + :returns: {str: bool} + A dictionary where keys are the names of submetrics and values are + whether a higher value of the submetric is better + """ + return { + "exact": True, # Exact match (the normalized answer exactly match the gold answer) + "f1": True, # The F-score of predicted tokens versus the gold answer + "HasAns_exact": True, # Exact match (the normalized answer exactly match the gold answer) + "HasAns_f1": True, # The F-score of predicted tokens versus the gold answer + "NoAns_exact": True, # Exact match (the normalized answer exactly match the gold answer) + "NoAns_f1": True, # The F-score of predicted tokens versus the gold answer + "best_exact": True, # Best exact match (with varying threshold) + "best_f1": True, # Best F1 (with varying threshold) + } diff --git a/kody/training_args.yaml b/kody/training_args.yaml deleted file mode 100644 index 52f1945..0000000 --- a/kody/training_args.yaml +++ /dev/null @@ -1,43 +0,0 @@ -cutoff_len: 1024 -dataset: alpaca_slovak_cleaned -dataset_dir: data -ddp_timeout: 180000000 -do_train: true -double_quantization: true -enable_thinking: false -eval_steps: 1000 -eval_strategy: steps -finetuning_type: lora -flash_attn: auto -fp16: true -gradient_accumulation_steps: 8 -include_num_input_tokens_seen: true -learning_rate: 0.0002 -logging_steps: 5 -lora_alpha: 32 -lora_dropout: 0.05 -lora_rank: 16 -lora_target: q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj -lr_scheduler_type: cosine -max_grad_norm: 1.0 -max_samples: 50000 -model_name_or_path: slovak-nlp/mistral-sk-7b -num_train_epochs: 1.0 -optim: adamw_8bit -output_dir: /home/schwarc/diplomovka/mistral_sk_alpaca/llamafactory-full-lora -packing: false -per_device_eval_batch_size: 1 -per_device_train_batch_size: 1 -plot_loss: true -preprocessing_num_workers: 16 -quantization_bit: 4 -quantization_method: bnb -report_to: none -save_steps: 1000 -save_total_limit: 2 -stage: sft -template: alpaca -trust_remote_code: true -use_unsloth: true -val_size: 0.025 -warmup_steps: 150 \ No newline at end of file diff --git a/plots/lmeval_main_metrics_chart.pdf b/plots/lmeval_main_metrics_chart.pdf new file mode 100644 index 0000000..6ad77b7 Binary files /dev/null and b/plots/lmeval_main_metrics_chart.pdf differ diff --git a/plots/lmeval_main_metrics_chart.png b/plots/lmeval_main_metrics_chart.png new file mode 100644 index 0000000..3197b9a Binary files /dev/null and b/plots/lmeval_main_metrics_chart.png differ diff --git a/results/lmeval_results_full_table.csv b/results/lmeval_results_full_table.csv new file mode 100644 index 0000000..b99bb75 --- /dev/null +++ b/results/lmeval_results_full_table.csv @@ -0,0 +1,18 @@ +task,samples,metric,value,model,adapter +arc_sk,1169,acc,0.3259,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +arc_sk,1169,acc_norm,0.3473,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +hellaswag_sk,9485,acc,0.4134,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +hellaswag_sk,9485,acc_norm,0.5106,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +m_mmlu_sk,13062,acc,0.2928,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +truthfulqa_sk_mc1,778,acc,0.2391,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +truthfulqa_sk_mc2,778,acc,0.3942,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +sklegal,1334,acc,0.2841,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +sklegal,1334,acc_norm,0.4693,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,exact,1.9722,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,f1,27.7994,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,HasAns_exact,2.4106,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,HasAns_f1,34.1456,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,NoAns_exact,0.0561,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,NoAns_f1,0.0561,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,best_exact,18.6163,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora +skquad,9583,best_f1,28.3012,slovak-nlp/mistral-sk-7b,Jakub1320/mistral-sk-7b-slovak-alpaca-qlora diff --git a/kody/evaluate_mistral_lora_metrics.py b/scripts/evaluate_mistral_lora_metrics.py similarity index 98% rename from kody/evaluate_mistral_lora_metrics.py rename to scripts/evaluate_mistral_lora_metrics.py index 562c12e..dddd0d8 100644 --- a/kody/evaluate_mistral_lora_metrics.py +++ b/scripts/evaluate_mistral_lora_metrics.py @@ -18,7 +18,6 @@ import sacrebleu MODEL_NAME = "slovak-nlp/mistral-sk-7b" DATASET_NAME = "saillab/alpaca-slovak-cleaned" -# Tvoj Unsloth + TRL + QLoRA adaptér ADAPTER_DIR = "/home/schwarc/diplomovka/mistral_sk_alpaca/mistral-sk-7b-alpaca-slovak-unsloth-lora-full" PROJECT_DIR = Path.home() / "diplomovka" / "evaluation_results" @@ -30,7 +29,7 @@ NUM_EVAL_SAMPLES = 1000 MAX_LENGTH = 1024 MAX_NEW_TOKENS = 300 -# Pri hodnotení je lepšie deterministické generovanie + DO_SAMPLE = False diff --git a/scripts/plot_lmeval_results.py b/scripts/plot_lmeval_results.py new file mode 100644 index 0000000..e698291 --- /dev/null +++ b/scripts/plot_lmeval_results.py @@ -0,0 +1,100 @@ +import pandas as pd +import matplotlib.pyplot as plt +from pathlib import Path + +OUT_DIR = Path("/home/schwarc/diplomovka/lmeval_plots") +OUT_DIR.mkdir(parents=True, exist_ok=True) + +MODEL = "slovak-nlp/mistral-sk-7b" +ADAPTER = "Jakub1320/mistral-sk-7b-slovak-alpaca-qlora" + +rows = [ + {"task": "arc_sk", "samples": 1169, "metric": "acc", "value": 0.3259}, + {"task": "arc_sk", "samples": 1169, "metric": "acc_norm", "value": 0.3473}, + + {"task": "hellaswag_sk", "samples": 9485, "metric": "acc", "value": 0.4134}, + {"task": "hellaswag_sk", "samples": 9485, "metric": "acc_norm", "value": 0.5106}, + + {"task": "m_mmlu_sk", "samples": 13062, "metric": "acc", "value": 0.2928}, + + {"task": "truthfulqa_sk_mc1", "samples": 778, "metric": "acc", "value": 0.2391}, + {"task": "truthfulqa_sk_mc2", "samples": 778, "metric": "acc", "value": 0.3942}, + + {"task": "sklegal", "samples": 1334, "metric": "acc", "value": 0.2841}, + {"task": "sklegal", "samples": 1334, "metric": "acc_norm", "value": 0.4693}, + + {"task": "skquad", "samples": 9583, "metric": "exact", "value": 1.9722}, + {"task": "skquad", "samples": 9583, "metric": "f1", "value": 27.7994}, + {"task": "skquad", "samples": 9583, "metric": "HasAns_exact", "value": 2.4106}, + {"task": "skquad", "samples": 9583, "metric": "HasAns_f1", "value": 34.1456}, + {"task": "skquad", "samples": 9583, "metric": "NoAns_exact", "value": 0.0561}, + {"task": "skquad", "samples": 9583, "metric": "NoAns_f1", "value": 0.0561}, + {"task": "skquad", "samples": 9583, "metric": "best_exact", "value": 18.6163}, + {"task": "skquad", "samples": 9583, "metric": "best_f1", "value": 28.3012}, +] + +df = pd.DataFrame(rows) +df["model"] = MODEL +df["adapter"] = ADAPTER + +csv_path = OUT_DIR / "lmeval_results_full_table.csv" +df.to_csv(csv_path, index=False) + +# Skrátený výber hlavných metrík pre graf + +main_rows = [ + {"benchmark": "ARC-SK", "metric": "acc_norm", "score_percent": 34.73}, + {"benchmark": "HellaSwag-SK", "metric": "acc_norm", "score_percent": 51.06}, + {"benchmark": "M-MMLU-SK", "metric": "acc", "score_percent": 29.28}, + {"benchmark": "TruthfulQA-SK MC1", "metric": "acc", "score_percent": 23.91}, + {"benchmark": "TruthfulQA-SK MC2", "metric": "acc", "score_percent": 39.42}, + {"benchmark": "SKLegal", "metric": "acc_norm", "score_percent": 46.93}, + {"benchmark": "SKQuAD", "metric": "f1", "score_percent": 27.80}, +] + +main_df = pd.DataFrame(main_rows) +main_csv_path = OUT_DIR / "lmeval_main_metrics_table.csv" +main_df.to_csv(main_csv_path, index=False) + +# Graf hlavných metrík +labels = main_df["benchmark"] + " (" + main_df["metric"] + ")" +scores = main_df["score_percent"] + +plt.figure(figsize=(11, 6)) +bars = plt.barh(labels, scores) + +plt.xlabel("Skóre (%)") +plt.ylabel("Benchmark") +plt.title("Výsledky lm-evaluation-harness pre PEFT LoRA model") +plt.xlim(0, 60) +plt.gca().invert_yaxis() + +for bar, score in zip(bars, scores): + plt.text( + bar.get_width() + 0.8, + bar.get_y() + bar.get_height() / 2, + f"{score:.2f} %", + va="center" + ) + +plt.figtext( + 0.01, + 0.01, + f"Model: {MODEL} | Adaptér: {ADAPTER}", + fontsize=8 +) + +plt.tight_layout(rect=[0, 0.04, 1, 1]) + +png_path = OUT_DIR / "lmeval_main_metrics_chart.png" +pdf_path = OUT_DIR / "lmeval_main_metrics_chart.pdf" + +plt.savefig(png_path, dpi=300) +plt.savefig(pdf_path) +plt.close() + +print("Hotovo.") +print(f"CSV kompletna tabulka: {csv_path}") +print(f"CSV hlavne metriky: {main_csv_path}") +print(f"Graf PNG: {png_path}") +print(f"Graf PDF: {pdf_path}") diff --git a/kody/test_mistral_lora.py b/scripts/test_mistral_lora.py similarity index 100% rename from kody/test_mistral_lora.py rename to scripts/test_mistral_lora.py diff --git a/kody/train_mistral_full.py b/scripts/train_mistral_full.py similarity index 100% rename from kody/train_mistral_full.py rename to scripts/train_mistral_full.py diff --git a/kody/train_mistral_unsloth_trl.py b/scripts/train_mistral_unsloth_trl.py similarity index 100% rename from kody/train_mistral_unsloth_trl.py rename to scripts/train_mistral_unsloth_trl.py diff --git a/strojovy_preklad/translate_dataset_to_slovak.py b/strojovy_preklad/translate_dataset_to_slovak.py new file mode 100644 index 0000000..3587c85 --- /dev/null +++ b/strojovy_preklad/translate_dataset_to_slovak.py @@ -0,0 +1,173 @@ +import json +from pathlib import Path + +import torch +from datasets import load_dataset +from transformers import AutoTokenizer, AutoModelForSeq2SeqLM + + +# ----------------------------- +# Nastavenia +# ----------------------------- + +DATASET_NAME = "tatsu-lab/alpaca" +SPLIT = "train" + +OUTPUT_DIR = Path("/home/schwarc/diplomovka/translated_datasets") +OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + +OUTPUT_FILE = OUTPUT_DIR / "alpaca_translated_en_sk_1000.jsonl" + +MAX_SAMPLES = 1000 + +TRANSLATION_MODEL = "facebook/nllb-200-distilled-600M" + +SOURCE_LANG = "eng_Latn" +TARGET_LANG = "slk_Latn" + +BATCH_SIZE = 8 +MAX_INPUT_LENGTH = 1024 +MAX_NEW_TOKENS = 512 + + + +# Pomocné funkcie + +def is_empty(value): + if value is None: + return True + value = str(value).strip() + return value == "" or value.lower() == "nan" + + +def load_already_done(path): + """ + Ak skript spadne alebo ho zastavíš, vie pokračovať. + Načíta už preložené riadky. + """ + if not path.exists(): + return [] + + rows = [] + with open(path, "r", encoding="utf-8") as f: + for line in f: + if line.strip(): + rows.append(json.loads(line)) + + return rows + + +def translate_batch(texts, tokenizer, model, device): + """ + Preloží batch textov z angličtiny do slovenčiny. + Prázdne texty nechá prázdne. + """ + results = [""] * len(texts) + + non_empty_indices = [] + non_empty_texts = [] + + for i, text in enumerate(texts): + if is_empty(text): + results[i] = "" + else: + non_empty_indices.append(i) + non_empty_texts.append(str(text).strip()) + + if not non_empty_texts: + return results + + tokenizer.src_lang = SOURCE_LANG + + inputs = tokenizer( + non_empty_texts, + return_tensors="pt", + padding=True, + truncation=True, + max_length=MAX_INPUT_LENGTH, + ).to(device) + + forced_bos_token_id = tokenizer.convert_tokens_to_ids(TARGET_LANG) + + with torch.no_grad(): + generated_tokens = model.generate( + **inputs, + forced_bos_token_id=forced_bos_token_id, + max_new_tokens=MAX_NEW_TOKENS, + num_beams=4, + ) + + translated = tokenizer.batch_decode( + generated_tokens, + skip_special_tokens=True, + ) + + for idx, translation in zip(non_empty_indices, translated): + results[idx] = translation.strip() + + return results + + +# Main + +def main(): + print("Loading dataset...") + dataset = load_dataset(DATASET_NAME, split=SPLIT) + + if MAX_SAMPLES is not None: + dataset = dataset.select(range(min(MAX_SAMPLES, len(dataset)))) + + print(f"Dataset size: {len(dataset)}") + + print("Loading translation model...") + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Device: {device}") + + tokenizer = AutoTokenizer.from_pretrained(TRANSLATION_MODEL) + model = AutoModelForSeq2SeqLM.from_pretrained( + TRANSLATION_MODEL, + torch_dtype=torch.float16 if device == "cuda" else torch.float32, + ).to(device) + + model.eval() + + already_done = load_already_done(OUTPUT_FILE) + start_idx = len(already_done) + + print(f"Already translated: {start_idx}") + print(f"Output file: {OUTPUT_FILE}") + + with open(OUTPUT_FILE, "a", encoding="utf-8") as out_f: + for start in range(start_idx, len(dataset), BATCH_SIZE): + end = min(start + BATCH_SIZE, len(dataset)) + batch = dataset[start:end] + + instructions_en = batch["instruction"] + inputs_en = batch["input"] + outputs_en = batch["output"] + + instructions_sk = translate_batch(instructions_en, tokenizer, model, device) + inputs_sk = translate_batch(inputs_en, tokenizer, model, device) + outputs_sk = translate_batch(outputs_en, tokenizer, model, device) + + for i in range(end - start): + row = { + "id": start + i, + "instruction_en": instructions_en[i], + "input_en": inputs_en[i], + "output_en": outputs_en[i], + "instruction_sk": instructions_sk[i], + "input_sk": inputs_sk[i], + "output_sk": outputs_sk[i], + } + + out_f.write(json.dumps(row, ensure_ascii=False) + "\n") + out_f.flush() + + print(f"Translated {end}/{len(dataset)}") + + print("Done.") + + +if __name__ == "__main__": + main()