|
|
|
|
|
|
|
|
import json |
|
|
from pathlib import Path |
|
|
from collections import Counter |
|
|
from sklearn.model_selection import train_test_split |
|
|
from datasets import Dataset, DatasetDict |
|
|
from transformers import AutoTokenizer |
|
|
|
|
|
|
|
|
INPUT_PATH = Path("data/labeled/labeled_dockerfiles.jsonl") |
|
|
TOP_RULES_PATH = Path("data/metadata/top_rules.json") |
|
|
OUTPUT_DIR = Path("data/processed/dataset_multilabel_top30") |
|
|
TOKENIZER_NAME = "microsoft/codebert-base" |
|
|
MAX_LENGTH = 512 |
|
|
SEED = 42 |
|
|
|
|
|
def load_top_rules(): |
|
|
with open(TOP_RULES_PATH, encoding="utf-8") as f: |
|
|
return json.load(f) |
|
|
|
|
|
def build_dataset(records, top_rules): |
|
|
rule2id = {r: i for i, r in enumerate(top_rules)} |
|
|
data = [] |
|
|
for row in records: |
|
|
if row.get("label") != "bad": |
|
|
continue |
|
|
|
|
|
triggered = row.get("rules_triggered", []) |
|
|
multilabel = [0] * len(top_rules) |
|
|
matched = False |
|
|
|
|
|
for rule in triggered: |
|
|
if rule in rule2id: |
|
|
multilabel[rule2id[rule]] = 1 |
|
|
matched = True |
|
|
|
|
|
if not matched: |
|
|
continue |
|
|
|
|
|
data.append({ |
|
|
"text": "\n".join(row["content"]) if isinstance(row["content"], list) else str(row["content"]), |
|
|
"labels": multilabel, |
|
|
"meta_lines": row.get("lines", {}), |
|
|
"meta_fixes": row.get("fixes", {}) |
|
|
}) |
|
|
|
|
|
return data |
|
|
|
|
|
def main(): |
|
|
print("📥 Wczytywanie danych...") |
|
|
top_rules = load_top_rules() |
|
|
print(f"🔝 Top {len(top_rules)} reguł: {top_rules}") |
|
|
|
|
|
with INPUT_PATH.open(encoding="utf-8") as f: |
|
|
records = [json.loads(line) for line in f if line.strip()] |
|
|
|
|
|
dataset = build_dataset(records, top_rules) |
|
|
print(f"📦 Zbudowano {len(dataset)} przykładów multilabel.") |
|
|
|
|
|
if not dataset: |
|
|
print("❌ Brak danych do przetworzenia. Sprawdź dane wejściowe.") |
|
|
return |
|
|
|
|
|
print("🔀 Podział na train/val/test...") |
|
|
train_val, test = train_test_split(dataset, test_size=0.1, random_state=SEED) |
|
|
train, val = train_test_split(train_val, test_size=0.1111, random_state=SEED) |
|
|
|
|
|
ds = DatasetDict({ |
|
|
"train": Dataset.from_list(train), |
|
|
"validation": Dataset.from_list(val), |
|
|
"test": Dataset.from_list(test), |
|
|
}) |
|
|
|
|
|
print("🔤 Tokenizacja...") |
|
|
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME) |
|
|
|
|
|
def tokenize_function(batch): |
|
|
texts = [str(x) if x is not None else "" for x in batch["text"]] |
|
|
return tokenizer( |
|
|
texts, |
|
|
padding="max_length", |
|
|
truncation=True, |
|
|
max_length=MAX_LENGTH |
|
|
) |
|
|
|
|
|
ds_tokenized = ds.map( |
|
|
tokenize_function, |
|
|
batched=True, |
|
|
remove_columns=["text", "meta_lines", "meta_fixes"] |
|
|
) |
|
|
|
|
|
print(f"💾 Zapisuję do: {OUTPUT_DIR}") |
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True) |
|
|
ds_tokenized.save_to_disk(str(OUTPUT_DIR)) |
|
|
|
|
|
print("✅ Gotowe.") |
|
|
|
|
|
if __name__ == "__main__": |
|
|
main() |
|
|
|
