dl_train.py

import os
import json
import random
import numpy as np
import torch
from pathlib import Path
from datasets import Dataset, ClassLabel
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, DataCollatorWithPadding
import evaluate
from sklearn.utils.class_weight import compute_class_weight

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed()

def load_code_dataset_from_multiple_folders(dataset_roots, file_exts=(".py", ".java", ".js")):
    data = []
    label2id = {}
    id2label = {}
    next_label_idx = 0
    for root in dataset_roots:
        if not os.path.exists(root):
            print(f"[WARNING] {root} not found, skipping.")
            continue
        print(f"[INFO] Processing {root}")
        folders = [f for f in os.listdir(root) if os.path.isdir(os.path.join(root, f))]
        if not folders:
            print(f"[WARNING] No subfolders in {root}, skipping.")
            continue
        for folder in sorted(folders):
            folder_path = os.path.join(root, folder)
            # Normalize labels to lowercase for consistency
            normalized_folder = folder.lower()
            if normalized_folder not in label2id:
                label2id[normalized_folder] = next_label_idx
                id2label[next_label_idx] = normalized_folder
                next_label_idx += 1
            files = []
            for ext in file_exts:
                files.extend(list(Path(folder_path).glob(f"*{ext}")))
            print(f"[INFO] Found {len(files)} files in {folder} -> {normalized_folder}")
            for file in files:
                try:
                    with open(file, "r", encoding="utf-8", errors="ignore") as f:
                        code = f.read()
                    if code.strip():
                        data.append({"text": code, "label": label2id[normalized_folder]})
                except Exception as e:
                    print(f"[WARNING] Could not read {file}: {e}")
    if not data:
        raise ValueError("No valid files found.")
    dataset = Dataset.from_list(data)
    names = [id2label[i] for i in range(len(id2label))]
    dataset = dataset.cast_column("label", ClassLabel(names=names))
    dataset = dataset.train_test_split(test_size=0.2, stratify_by_column="label", seed=42)
    dataset["validation"] = dataset["train"].train_test_split(test_size=0.1, stratify_by_column="label", seed=42)["test"]
    return dataset, label2id, id2label

def compute_metrics(eval_pred):
    acc_metric = evaluate.load("accuracy")
    f1_metric = evaluate.load("f1")
    logits, labels = eval_pred
    if isinstance(logits, tuple):
        logits = logits[0]
    predictions = np.argmax(logits, axis=-1)
    labels = np.array(labels)
    acc = acc_metric.compute(predictions=predictions, references=labels)
    f1_macro = f1_metric.compute(predictions=predictions, references=labels, average="macro")
    f1_weighted = f1_metric.compute(predictions=predictions, references=labels, average="weighted")
    return {**acc, "f1_macro": f1_macro["f1"], "f1_weighted": f1_weighted["f1"]}

class WeightedTrainer(Trainer):
    def __init__(self, class_weights=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.class_weights = class_weights

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        if self.class_weights is not None:
            cw = self.class_weights.to(logits.device)
            loss_fct = torch.nn.CrossEntropyLoss(weight=cw)
            loss = loss_fct(logits.view(-1, model.config.num_labels), labels.view(-1))
        else:
            loss = outputs.loss
        return (loss, outputs) if return_outputs else loss

def train_and_evaluate_model(model_name, dataset, label2id, id2label, output_dir, epochs=5, batch_size=8, lr=1e-5):
    print(f"[INFO] Training {model_name}")
    os.makedirs(output_dir, exist_ok=True)
    labels = dataset["train"]["label"]
    weights = compute_class_weight("balanced", classes=np.unique(labels), y=np.array(labels))
    weight_vector = np.ones(len(label2id))
    for idx, w in zip(np.unique(labels), weights):
        weight_vector[idx] = w
    class_weights = torch.FloatTensor(weight_vector)
    print(f"[INFO] Class weights computed.")

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenized = dataset.map(lambda x: tokenizer(x["text"], truncation=True, max_length=512), batched=True, remove_columns=["text"])
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(label2id), id2label=id2label, label2id=label2id)

    training_args = TrainingArguments(
        output_dir=output_dir,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        learning_rate=lr,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        num_train_epochs=epochs,
        weight_decay=0.01,
        load_best_model_at_end=True,
        metric_for_best_model="f1_macro",
        logging_dir=os.path.join(output_dir, "logs"),
        save_total_limit=2,
        logging_steps=20,
    )
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    trainer = WeightedTrainer(
        class_weights=class_weights,
        model=model,
        args=training_args,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized["validation"],
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    trainer.train()
    print("[INFO] Evaluating...")
    results = trainer.evaluate(tokenized["test"])
    print(f"[RESULT] Accuracy={results['eval_accuracy']:.4f}, F1 Macro={results['eval_f1_macro']:.4f}, F1 Weighted={results['eval_f1_weighted']:.4f}")
    trainer.save_model(output_dir)
    with open(os.path.join(output_dir, "label_mappings.json"), "w") as f:
        json.dump({"label2id": label2id, "id2label": id2label}, f, indent=2)
    return results

def main():
    DATASET_ROOTS = [
        r"D:\Projects\Python\Code_Detector\Dataset\Python",
        r"D:\Projects\Python\Code_Detector\Dataset\Java",
        r"D:\Projects\Python\Code_Detector\Dataset\JS"
    ]
    MODELS = {
        "CodeBERT": "microsoft/codebert-base",
        "GraphCodeBERT": "microsoft/graphcodebert-base"
    }
    OUTPUT_BASE = r"D:\Projects\Python\Code_Detector\output"
    EPOCHS = 5
    BATCH_SIZE = 8
    LR = 1e-5

    print("[INFO] Loading dataset...")
    dataset, label2id, id2label = load_code_dataset_from_multiple_folders(DATASET_ROOTS)
    print("[INFO] Dataset loaded.")
    for model_label, model_name in MODELS.items():
        output_dir = os.path.join(OUTPUT_BASE, model_label)
        train_and_evaluate_model(model_name, dataset, label2id, id2label, output_dir, epochs=EPOCHS, batch_size=BATCH_SIZE, lr=LR)

if __name__ == "__main__":
    main()