| | |
| |
|
| | 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() |
| |
|
