From 26dcca5ebd07e1ac5769078831457a4330b84147 Mon Sep 17 00:00:00 2001 From: HardikChhallani <116100549+HardikChhallani@users.noreply.github.com> Date: Tue, 9 Sep 2025 23:25:29 +0530 Subject: [PATCH 1/2] Updated Models and finetuning --- .idea/workspace.xml | 14 ++- Processing/finetuning/adfs_qwen2.py | 13 +-- Processing/finetuning/base_finetuning.py | 110 +++++++++++++++++---- Processing/finetuning/hdfs_qwen2.py | 17 ++-- Processing/finetuning/openstack_qwen2.py | 18 ++-- Processing/finetuning/thunderbird_qwen2.py | 18 ++-- main.py | 16 --- template.py | 0 8 files changed, 129 insertions(+), 77 deletions(-) delete mode 100644 main.py delete mode 100644 template.py diff --git a/.idea/workspace.xml b/.idea/workspace.xml index 692d7fc..f6eeb0a 100644 --- a/.idea/workspace.xml +++ b/.idea/workspace.xml @@ -6,7 +6,11 @@ + + + + diff --git a/Processing/finetuning/adfs_qwen2.py b/Processing/finetuning/adfs_qwen2.py index 22dbbec..7842ed2 100644 --- a/Processing/finetuning/adfs_qwen2.py +++ b/Processing/finetuning/adfs_qwen2.py @@ -4,20 +4,17 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOCAL_MODEL_PATH = "models/adfa_qwen2" +LOG_NAME = "adfs" class ADFAQwen2FineTune(BaseFineTune): def __init__(self): - self.model_name = MODEL_NAME - super().__init__(self.model_name,logger=logger) + super().__init__(MODEL_NAME, logger=logger) - def finetune(self): - # Implement fine-tuning logic for ADFA Qwen2 model - pass + def run_finetuning(self, dataset): + self.finetune(LOG_NAME, dataset) def save_checkpoint(self): - # Implement saving checkpoint logic for ADFA Qwen2 model - pass + self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") def upload_to_hf(self): # Implement uploading to Hugging Face logic for ADFA Qwen2 model diff --git a/Processing/finetuning/base_finetuning.py b/Processing/finetuning/base_finetuning.py index ed272c0..4343c95 100644 --- a/Processing/finetuning/base_finetuning.py +++ b/Processing/finetuning/base_finetuning.py @@ -1,12 +1,13 @@ import logging import os from abc import abstractmethod, ABC -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer +from peft import LoraConfig, get_peft_model -LOCAL_MODEL_PATH = "QWEN2/qwen2.5-1.5B-instruct" +BASE_MODEL_PATH = "QWEN2/qwen2.5-1.5B-instruct" class BaseFineTune(ABC): - def __int__(self, + def __init__(self, model_name: str, logger = None, *args, **kwargs): @@ -15,37 +16,106 @@ def __int__(self, self.model = None self.logger = logger if logger else logging.getLogger(__name__) - if not os.path.exists(LOCAL_MODEL_PATH): - self.load() + if not os.path.exists(BASE_MODEL_PATH): + self.download_model() + self.load_model() - def load(self): + def download_model(self): """ - Load the ADFA Qwen2 model and tokenizer. + Download the model and tokenizer from Hugging Face. """ try: - self.logger.info("Loading ADFA Qwen2 model...") - self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) - self.model = AutoModelForCausalLM.from_pretrained( + self.logger.info(f"Downloading model {self.model_name}...") + tokenizer = AutoTokenizer.from_pretrained(self.model_name) + model = AutoModelForCausalLM.from_pretrained( self.model_name, torch_dtype="auto", device_map="auto" ) - self.logger.info("Model and tokenizer loaded successfully.") - self.logger.info(f"Saving model and tokenizer to: {self.local_model_path}") - self.model.save_pretrained(LOCAL_MODEL_PATH) - self.tokenizer.save_pretrained(LOCAL_MODEL_PATH) + self.logger.info("Model and tokenizer downloaded successfully.") + self.logger.info(f"Saving model and tokenizer to: {BASE_MODEL_PATH}") + model.save_pretrained(BASE_MODEL_PATH) + tokenizer.save_pretrained(BASE_MODEL_PATH) self.logger.info("✅ Model saved successfully to your local machine!") except Exception as e: - self.logger.error(f"Error loading model {self.model_name}: {e}") + self.logger.error(f"Error downloading model {self.model_name}: {e}") raise e - @abstractmethod - def finetune(self): - pass + def load_model(self): + """ + Load the model and tokenizer from the local directory. + """ + try: + self.logger.info(f"Loading model from {BASE_MODEL_PATH}...") + self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH) + self.model = AutoModelForCausalLM.from_pretrained( + BASE_MODEL_PATH, + torch_dtype="auto", + device_map="auto" + ) + self.logger.info("Model and tokenizer loaded successfully.") + except Exception as e: + self.logger.error(f"Error loading model from {BASE_MODEL_PATH}: {e}") + raise e + + def finetune(self, log_name: str, dataset): + """ + Fine-tune the model and save the fine-tuned weights. + """ + try: + self.logger.info(f"Starting fine-tuning for {log_name}...") + + # Set up LoRA configuration + lora_config = LoraConfig( + r=8, + lora_alpha=32, + lora_dropout=0.1, + target_modules=["q_proj", "v_proj"] + ) + self.model = get_peft_model(self.model, lora_config) + + # Set up training arguments + training_args = TrainingArguments( + output_dir=f"./{log_name}_model", + per_device_train_batch_size=4, + gradient_accumulation_steps=4, + learning_rate=2e-4, + num_train_epochs=3, + logging_steps=10, + save_steps=100, + evaluation_strategy="steps", + eval_steps=100, + load_best_model_at_end=True, + ) + + # Create Trainer + self.trainer = Trainer( + model=self.model, + args=training_args, + train_dataset=dataset, + eval_dataset=dataset, + ) + + # Start fine-tuning + self.trainer.train() + + # Save the fine-tuned model + self.logger.info(f"Saving fine-tuned model to ./{log_name}_model") + self.trainer.save_model() + self.logger.info("✅ Fine-tuned model saved successfully!") + + except Exception as e: + self.logger.error(f"Error during fine-tuning for {log_name}: {e}") + raise e @abstractmethod - def save_checkpoint(self): - pass + def save_checkpoint(self, checkpoint_path: str): + if self.trainer: + self.logger.info(f"Saving checkpoint to {checkpoint_path}") + self.trainer.save_model(checkpoint_path) + self.logger.info("✅ Checkpoint saved successfully!") + else: + self.logger.warning("Trainer not initialized. Cannot save checkpoint.") @abstractmethod def upload_to_hf(self): diff --git a/Processing/finetuning/hdfs_qwen2.py b/Processing/finetuning/hdfs_qwen2.py index ba624b7..3501425 100644 --- a/Processing/finetuning/hdfs_qwen2.py +++ b/Processing/finetuning/hdfs_qwen2.py @@ -4,21 +4,18 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOCAL_MODEL_PATH = "models/hdfs_qwen2" +LOG_NAME = "hdfs" -class ADFAQwen2FineTune(BaseFineTune): +class HDFSQwen2FineTune(BaseFineTune): def __init__(self): - self.model_name = MODEL_NAME - super().__init__(self.model_name,logger=logger) + super().__init__(MODEL_NAME, logger=logger) - def finetune(self): - # Implement fine-tuning logic for ADFA Qwen2 model - pass + def run_finetuning(self, dataset): + self.finetune(LOG_NAME, dataset) def save_checkpoint(self): - # Implement saving checkpoint logic for ADFA Qwen2 model - pass + self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") def upload_to_hf(self): - # Implement uploading to Hugging Face logic for ADFA Qwen2 model + # Implement uploading to Hugging Face logic for HDFS Qwen2 model pass \ No newline at end of file diff --git a/Processing/finetuning/openstack_qwen2.py b/Processing/finetuning/openstack_qwen2.py index c020603..77220dc 100644 --- a/Processing/finetuning/openstack_qwen2.py +++ b/Processing/finetuning/openstack_qwen2.py @@ -4,22 +4,18 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOCAL_MODEL_PATH = "models/openstack_qwen2" +LOG_NAME = "openstack" -class ADFAQwen2FineTune(BaseFineTune): +class OpenStackQwen2FineTune(BaseFineTune): def __init__(self): - self.model_name = MODEL_NAME - self.local_model_path = LOCAL_MODEL_PATH - super().__init__(self.model_name,logger=logger) + super().__init__(MODEL_NAME, logger=logger) - def finetune(self): - # Implement fine-tuning logic for ADFA Qwen2 model - pass + def run_finetuning(self, dataset): + self.finetune(LOG_NAME, dataset) def save_checkpoint(self): - # Implement saving checkpoint logic for ADFA Qwen2 model - pass + self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") def upload_to_hf(self): - # Implement uploading to Hugging Face logic for ADFA Qwen2 model + # Implement uploading to Hugging Face logic for OpenStack Qwen2 model pass \ No newline at end of file diff --git a/Processing/finetuning/thunderbird_qwen2.py b/Processing/finetuning/thunderbird_qwen2.py index 6b2683d..a93563f 100644 --- a/Processing/finetuning/thunderbird_qwen2.py +++ b/Processing/finetuning/thunderbird_qwen2.py @@ -4,22 +4,18 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOCAL_MODEL_PATH = "models/thunderbird_qwen2" +LOG_NAME = "thunderbird" -class ADFAQwen2FineTune(BaseFineTune): +class ThunderbirdQwen2FineTune(BaseFineTune): def __init__(self): - self.model_name = MODEL_NAME - self.local_model_path = LOCAL_MODEL_PATH - super().__init__(self.model_name,logger=logger) + super().__init__(MODEL_NAME, logger=logger) - def finetune(self): - # Implement fine-tuning logic for ADFA Qwen2 model - pass + def run_finetuning(self, dataset): + self.finetune(LOG_NAME, dataset) def save_checkpoint(self): - # Implement saving checkpoint logic for ADFA Qwen2 model - pass + self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") def upload_to_hf(self): - # Implement uploading to Hugging Face logic for ADFA Qwen2 model + # Implement uploading to Hugging Face logic for Thunderbird Qwen2 model pass \ No newline at end of file diff --git a/main.py b/main.py deleted file mode 100644 index 94e3a87..0000000 --- a/main.py +++ /dev/null @@ -1,16 +0,0 @@ -# This is a sample Python script. - -# Press ⌃R to execute it or replace it with your code. -# Press Double ⇧ to search everywhere for classes, files, tool windows, actions, and settings. - - -def print_hi(name): - # Use a breakpoint in the code line below to debug your script. - print(f'Hi, {name}') # Press ⌘F8 to toggle the breakpoint. - - -# Press the green button in the gutter to run the script. -if __name__ == '__main__': - print_hi('PyCharm') - -# See PyCharm help at https://www.jetbrains.com/help/pycharm/ diff --git a/template.py b/template.py deleted file mode 100644 index e69de29..0000000 From acd68e8004c16b57fcd9eb2fa37ba0e44ff4d930 Mon Sep 17 00:00:00 2001 From: HardikChhallani <116100549+HardikChhallani@users.noreply.github.com> Date: Tue, 4 Nov 2025 22:44:53 +0530 Subject: [PATCH 2/2] Updated Models and finetuning --- .idea/workspace.xml | 55 +++++---- Dataset/template.py | 7 -- Processing/Dataset/base_dataset.py | 133 +++++++++++++++------ Processing/finetuning/adfs_qwen2.py | 126 ++++++++++++++++--- Processing/finetuning/base_finetuning.py | 122 ------------------- Processing/finetuning/hdfs_qwen2.py | 126 ++++++++++++++++--- Processing/finetuning/openstack_qwen2.py | 126 ++++++++++++++++--- Processing/finetuning/thunderbird_qwen2.py | 126 ++++++++++++++++--- src/SystemLogLLM/main.py | 126 ++++++++++++++++++- 9 files changed, 691 insertions(+), 256 deletions(-) delete mode 100644 Dataset/template.py delete mode 100644 Processing/finetuning/base_finetuning.py diff --git a/.idea/workspace.xml b/.idea/workspace.xml index f6eeb0a..afb7e61 100644 --- a/.idea/workspace.xml +++ b/.idea/workspace.xml @@ -4,13 +4,16 @@ - + + + - + + @@ -58,19 +62,28 @@ "associatedIndex": 0 } + + + - { - "keyToString": { - "ModuleVcsDetector.initialDetectionPerformed": "true", - "RunOnceActivity.ShowReadmeOnStart": "true", - "RunOnceActivity.git.unshallow": "true", - "git-widget-placeholder": "Hardik/Finetune" + +}]]> + + + + + @@ -163,7 +176,15 @@ @@ -172,17 +193,7 @@ - - - - - - file://$PROJECT_DIR$/main.py - 8 - - - + + \ No newline at end of file diff --git a/Dataset/template.py b/Dataset/template.py deleted file mode 100644 index 9749066..0000000 --- a/Dataset/template.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -Save all dataset in this path with template - -Dataset Name: / - Raw Dataset - Test - Train -""" \ No newline at end of file diff --git a/Processing/Dataset/base_dataset.py b/Processing/Dataset/base_dataset.py index 7144398..a25d938 100644 --- a/Processing/Dataset/base_dataset.py +++ b/Processing/Dataset/base_dataset.py @@ -1,38 +1,95 @@ -from abc import ABC, abstractmethod - -class BaseDataset(ABC): - """ - Base class for datasets. - - Datasets should inherit from this class and implement the required methods - for preprocessing and create required data structures for training. - """ - - def __init__(self, - dataset_name: str, - dataset_path: str, - **kwargs): - """ - Initialize the dataset. - - Args: - dataset_name: Name of the dataset - split: Split of the dataset - **kwargs: Additional dataset-specific arguments - """ - self.dataset_name = dataset_name - self.dataset_path = dataset_path - - def load_csv(self): - """ - Load dataset from a CSV file. - - Returns: - Loaded dataset - """ - import pandas as pd - return pd.read_csv(self.dataset_path) - - @abstractmethod - def format_csv(self): - pass +import os, json, random, requests, tarfile +from io import BytesIO +from google.colab import drive + +drive.mount('/content/drive') +drive_dir = "/content/drive/MyDrive/adfa_finetune_data" +os.makedirs(drive_dir, exist_ok=True) +print(f"🔹 All datasets will be saved in: {drive_dir}") + +datasets = { + "adfa": { + "normal": "https://raw.githubusercontent.com/ait-aecid/anomaly-detection-log-datasets/refs/heads/main/adfa_verazuo/adfa_test_normal", + "abnormal": "https://raw.githubusercontent.com/ait-aecid/anomaly-detection-log-datasets/refs/heads/main/adfa_verazuo/adfa_test_abnormal" + }, + "hdfs": { + "normal": "https://raw.githubusercontent.com/ait-aecid/anomaly-detection-log-datasets/refs/heads/main/hdfs_loghub/hdfs_test_normal", + "abnormal": "https://raw.githubusercontent.com/ait-aecid/anomaly-detection-log-datasets/refs/heads/main/hdfs_loghub/hdfs_test_abnormal" + }, + "openstack": { + "normal": "https://raw.githubusercontent.com/ait-aecid/anomaly-detection-log-datasets/refs/heads/main/openstack_loghub/openstack_test_normal", + "abnormal": "https://raw.githubusercontent.com/ait-aecid/anomaly-detection-log-datasets/refs/heads/main/openstack_loghub/openstack_test_abnormal" + }, + "thunderbird": { + "normal": "https://github.com/ait-aecid/anomaly-detection-log-datasets/raw/main/thunderbird_cfdr/thunderbird_test_normal.tar.gz", + "abnormal": "https://github.com/ait-aecid/anomaly-detection-log-datasets/raw/main/thunderbird_cfdr/thunderbird_test_abnormal.tar.gz" + } +} + +def read_plain(url): + print(f"📥 Downloading {url}") + r = requests.get(url) + r.raise_for_status() + return [l.strip() for l in r.text.splitlines() if l.strip()] + +def read_tar(url): + print(f"📦 Extracting {url}") + r = requests.get(url) + r.raise_for_status() + lines=[] + with tarfile.open(fileobj=BytesIO(r.content), mode="r:gz") as tar: + for m in tar.getmembers(): + if m.isfile(): + f = tar.extractfile(m) + if f: + txt = f.read().decode(errors="ignore") + lines += [l.strip() for l in txt.splitlines() if l.strip()] + return lines + +os.makedirs("data", exist_ok=True) + +for name, urls in datasets.items(): + print(f"\n🔹 Processing dataset: {name}") + if name == "thunderbird": + normals = read_tar(urls["normal"]) + abnormals = read_tar(urls["abnormal"]) + else: + normals = read_plain(urls["normal"]) + abnormals = read_plain(urls["abnormal"]) + + all_items = ( + [{"dataset": name, "text": t, "label": "normal"} for t in normals] + + [{"dataset": name, "text": t, "label": "abnormal"} for t in abnormals] + ) + random.shuffle(all_items) + + split = int(0.8 * len(all_items)) + train = all_items[:split] + test = all_items[split:] + + def fmt(e): + return { + "instruction": "Classify the following system log as normal or abnormal.", + "input": f"[{e['dataset'].upper()}] {e['text']}", + "output": e["label"] + } + + train_fmt = [fmt(e) for e in train] + test_fmt = [fmt(e) for e in test] + + local_train = f"data/train_{name}.json" + local_test = f"data/test_{name}.json" + drive_train = os.path.join(drive_dir, f"train_{name}.json") + drive_test = os.path.join(drive_dir, f"test_{name}.json") + + for path, data in [(local_train, train_fmt), (local_test, test_fmt), + (drive_train, train_fmt), (drive_test, test_fmt)]: + with open(path, "w", encoding="utf-8") as f: + for o in data: + f.write(json.dumps(o, ensure_ascii=False) + "\n") + + print(f"✅ Saved {len(train_fmt)} train + {len(test_fmt)} test entries for {name}") + print(f" ↳ {drive_train}") + print(f" ↳ {drive_test}") + +print("\n🎯 All datasets processed and saved to Drive!") \ No newline at end of file diff --git a/Processing/finetuning/adfs_qwen2.py b/Processing/finetuning/adfs_qwen2.py index 7842ed2..19dcffe 100644 --- a/Processing/finetuning/adfs_qwen2.py +++ b/Processing/finetuning/adfs_qwen2.py @@ -1,21 +1,115 @@ -from base_finetuning import BaseFineTune -import logging +import os +import torch +from datasets import load_dataset +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + TrainingArguments, + Trainer, +) +from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training +from transformers import BitsAndBytesConfig -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOG_NAME = "adfs" +train_path = "/content/data/train_adfa.json" +test_path = "/content/data/test_adfa.json" -class ADFAQwen2FineTune(BaseFineTune): - def __init__(self): - super().__init__(MODEL_NAME, logger=logger) +dataset_name = os.path.splitext(os.path.basename(train_path))[0].replace("train_", "") +save_dir = f"./qwen_lora_{dataset_name}_model" - def run_finetuning(self, dataset): - self.finetune(LOG_NAME, dataset) +print("🧠 Loading datasets...") +dataset = load_dataset( + "json", + data_files={"train": train_path, "validation": test_path}, +) - def save_checkpoint(self): - self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") +print(dataset) - def upload_to_hf(self): - # Implement uploading to Hugging Face logic for ADFA Qwen2 model - pass \ No newline at end of file +def build_prompt(instruction, inp, output): + """Format each sample into an instruction-style prompt.""" + return ( + f"Instruction: {instruction}\n" + f"Input: {inp}\n" + f"Response: {output}" + ) + +print("🔡 Loading tokenizer...") +model_name = "Qwen/Qwen2-1.5B-Instruct" +tokenizer = AutoTokenizer.from_pretrained(model_name) +tokenizer.pad_token = tokenizer.eos_token + +bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + bnb_4bit_compute_dtype=torch.bfloat16, +) + +print("⚙️ Loading model with 4-bit quantization...") +model = AutoModelForCausalLM.from_pretrained( + model_name, + quantization_config=bnb_config, + device_map="auto", +) +model = prepare_model_for_kbit_training(model) + +print("🪶 Applying LoRA configuration...") +lora_config = LoraConfig( + r=64, + lora_alpha=16, + target_modules=["q_proj", "v_proj"], + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", +) +model = get_peft_model(model, lora_config) +model.print_trainable_parameters() + +print("✂️ Tokenizing dataset...") + +def tokenize_function(batch): + texts = [ + build_prompt(inst, inp, out) + for inst, inp, out in zip(batch["instruction"], batch["input"], batch["output"]) + ] + tokens = tokenizer( + texts, + padding="max_length", + truncation=True, + max_length=1024, + ) + tokens["labels"] = tokens["input_ids"].copy() + return tokens + +tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset["train"].column_names) + +args = TrainingArguments( + output_dir=save_dir, + per_device_train_batch_size=1, + gradient_accumulation_steps=4, + num_train_epochs=1, + learning_rate=2e-4, + fp16=True, + save_total_limit=2, + evaluation_strategy="epoch", + logging_dir="./logs", + logging_steps=25, + report_to="none", +) + +trainer = Trainer( + model=model, + args=args, + train_dataset=tokenized_dataset["train"], + eval_dataset=tokenized_dataset["validation"], + tokenizer=tokenizer, +) + +print("🚀 Starting training...") +trainer.train() + +print(f"💾 Saving LoRA fine-tuned model to: {save_dir}") +os.makedirs(save_dir, exist_ok=True) +model.save_pretrained(save_dir) +tokenizer.save_pretrained(save_dir) + +print(f"✅ Model and tokenizer saved successfully at: {save_dir}") \ No newline at end of file diff --git a/Processing/finetuning/base_finetuning.py b/Processing/finetuning/base_finetuning.py deleted file mode 100644 index 4343c95..0000000 --- a/Processing/finetuning/base_finetuning.py +++ /dev/null @@ -1,122 +0,0 @@ -import logging -import os -from abc import abstractmethod, ABC -from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer -from peft import LoraConfig, get_peft_model - -BASE_MODEL_PATH = "QWEN2/qwen2.5-1.5B-instruct" - -class BaseFineTune(ABC): - def __init__(self, - model_name: str, - logger = None, - *args, **kwargs): - self.model_name = model_name - self.tokenizer = None - self.model = None - self.logger = logger if logger else logging.getLogger(__name__) - - if not os.path.exists(BASE_MODEL_PATH): - self.download_model() - self.load_model() - - def download_model(self): - """ - Download the model and tokenizer from Hugging Face. - """ - try: - self.logger.info(f"Downloading model {self.model_name}...") - tokenizer = AutoTokenizer.from_pretrained(self.model_name) - model = AutoModelForCausalLM.from_pretrained( - self.model_name, - torch_dtype="auto", - device_map="auto" - ) - self.logger.info("Model and tokenizer downloaded successfully.") - self.logger.info(f"Saving model and tokenizer to: {BASE_MODEL_PATH}") - model.save_pretrained(BASE_MODEL_PATH) - tokenizer.save_pretrained(BASE_MODEL_PATH) - self.logger.info("✅ Model saved successfully to your local machine!") - except Exception as e: - self.logger.error(f"Error downloading model {self.model_name}: {e}") - raise e - - def load_model(self): - """ - Load the model and tokenizer from the local directory. - """ - try: - self.logger.info(f"Loading model from {BASE_MODEL_PATH}...") - self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_PATH) - self.model = AutoModelForCausalLM.from_pretrained( - BASE_MODEL_PATH, - torch_dtype="auto", - device_map="auto" - ) - self.logger.info("Model and tokenizer loaded successfully.") - except Exception as e: - self.logger.error(f"Error loading model from {BASE_MODEL_PATH}: {e}") - raise e - - def finetune(self, log_name: str, dataset): - """ - Fine-tune the model and save the fine-tuned weights. - """ - try: - self.logger.info(f"Starting fine-tuning for {log_name}...") - - # Set up LoRA configuration - lora_config = LoraConfig( - r=8, - lora_alpha=32, - lora_dropout=0.1, - target_modules=["q_proj", "v_proj"] - ) - self.model = get_peft_model(self.model, lora_config) - - # Set up training arguments - training_args = TrainingArguments( - output_dir=f"./{log_name}_model", - per_device_train_batch_size=4, - gradient_accumulation_steps=4, - learning_rate=2e-4, - num_train_epochs=3, - logging_steps=10, - save_steps=100, - evaluation_strategy="steps", - eval_steps=100, - load_best_model_at_end=True, - ) - - # Create Trainer - self.trainer = Trainer( - model=self.model, - args=training_args, - train_dataset=dataset, - eval_dataset=dataset, - ) - - # Start fine-tuning - self.trainer.train() - - # Save the fine-tuned model - self.logger.info(f"Saving fine-tuned model to ./{log_name}_model") - self.trainer.save_model() - self.logger.info("✅ Fine-tuned model saved successfully!") - - except Exception as e: - self.logger.error(f"Error during fine-tuning for {log_name}: {e}") - raise e - - @abstractmethod - def save_checkpoint(self, checkpoint_path: str): - if self.trainer: - self.logger.info(f"Saving checkpoint to {checkpoint_path}") - self.trainer.save_model(checkpoint_path) - self.logger.info("✅ Checkpoint saved successfully!") - else: - self.logger.warning("Trainer not initialized. Cannot save checkpoint.") - - @abstractmethod - def upload_to_hf(self): - pass \ No newline at end of file diff --git a/Processing/finetuning/hdfs_qwen2.py b/Processing/finetuning/hdfs_qwen2.py index 3501425..bf34d66 100644 --- a/Processing/finetuning/hdfs_qwen2.py +++ b/Processing/finetuning/hdfs_qwen2.py @@ -1,21 +1,115 @@ -from base_finetuning import BaseFineTune -import logging +import os +import torch +from datasets import load_dataset +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + TrainingArguments, + Trainer, +) +from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training +from transformers import BitsAndBytesConfig -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOG_NAME = "hdfs" +train_path = "/content/data/train_hdfs.json" +test_path = "/content/data/test_hdfs.json" -class HDFSQwen2FineTune(BaseFineTune): - def __init__(self): - super().__init__(MODEL_NAME, logger=logger) +dataset_name = os.path.splitext(os.path.basename(train_path))[0].replace("train_", "") +save_dir = f"./qwen_lora_{dataset_name}_model" - def run_finetuning(self, dataset): - self.finetune(LOG_NAME, dataset) +print("🧠 Loading datasets...") +dataset = load_dataset( + "json", + data_files={"train": train_path, "validation": test_path}, +) - def save_checkpoint(self): - self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") +print(dataset) - def upload_to_hf(self): - # Implement uploading to Hugging Face logic for HDFS Qwen2 model - pass \ No newline at end of file +def build_prompt(instruction, inp, output): + """Format each sample into an instruction-style prompt.""" + return ( + f"Instruction: {instruction}\n" + f"Input: {inp}\n" + f"Response: {output}" + ) + +print("🔡 Loading tokenizer...") +model_name = "Qwen/Qwen2-1.5B-Instruct" +tokenizer = AutoTokenizer.from_pretrained(model_name) +tokenizer.pad_token = tokenizer.eos_token + +bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + bnb_4bit_compute_dtype=torch.bfloat16, +) + +print("⚙️ Loading model with 4-bit quantization...") +model = AutoModelForCausalLM.from_pretrained( + model_name, + quantization_config=bnb_config, + device_map="auto", +) +model = prepare_model_for_kbit_training(model) + +print("🪶 Applying LoRA configuration...") +lora_config = LoraConfig( + r=64, + lora_alpha=16, + target_modules=["q_proj", "v_proj"], + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", +) +model = get_peft_model(model, lora_config) +model.print_trainable_parameters() + +print("✂️ Tokenizing dataset...") + +def tokenize_function(batch): + texts = [ + build_prompt(inst, inp, out) + for inst, inp, out in zip(batch["instruction"], batch["input"], batch["output"]) + ] + tokens = tokenizer( + texts, + padding="max_length", + truncation=True, + max_length=1024, + ) + tokens["labels"] = tokens["input_ids"].copy() + return tokens + +tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset["train"].column_names) + +args = TrainingArguments( + output_dir=save_dir, + per_device_train_batch_size=1, + gradient_accumulation_steps=4, + num_train_epochs=1, + learning_rate=2e-4, + fp16=True, + save_total_limit=2, + evaluation_strategy="epoch", + logging_dir="./logs", + logging_steps=25, + report_to="none", +) + +trainer = Trainer( + model=model, + args=args, + train_dataset=tokenized_dataset["train"], + eval_dataset=tokenized_dataset["validation"], + tokenizer=tokenizer, +) + +print("🚀 Starting training...") +trainer.train() + +print(f"💾 Saving LoRA fine-tuned model to: {save_dir}") +os.makedirs(save_dir, exist_ok=True) +model.save_pretrained(save_dir) +tokenizer.save_pretrained(save_dir) + +print(f"✅ Model and tokenizer saved successfully at: {save_dir}") \ No newline at end of file diff --git a/Processing/finetuning/openstack_qwen2.py b/Processing/finetuning/openstack_qwen2.py index 77220dc..6c59ec0 100644 --- a/Processing/finetuning/openstack_qwen2.py +++ b/Processing/finetuning/openstack_qwen2.py @@ -1,21 +1,115 @@ -from base_finetuning import BaseFineTune -import logging +import os +import torch +from datasets import load_dataset +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + TrainingArguments, + Trainer, +) +from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training +from transformers import BitsAndBytesConfig -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOG_NAME = "openstack" +train_path = "/content/data/train_openstack.json" +test_path = "/content/data/test_openstack.json" -class OpenStackQwen2FineTune(BaseFineTune): - def __init__(self): - super().__init__(MODEL_NAME, logger=logger) +dataset_name = os.path.splitext(os.path.basename(train_path))[0].replace("train_", "") +save_dir = f"./qwen_lora_{dataset_name}_model" - def run_finetuning(self, dataset): - self.finetune(LOG_NAME, dataset) +print("🧠 Loading datasets...") +dataset = load_dataset( + "json", + data_files={"train": train_path, "validation": test_path}, +) - def save_checkpoint(self): - self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") +print(dataset) - def upload_to_hf(self): - # Implement uploading to Hugging Face logic for OpenStack Qwen2 model - pass \ No newline at end of file +def build_prompt(instruction, inp, output): + """Format each sample into an instruction-style prompt.""" + return ( + f"Instruction: {instruction}\n" + f"Input: {inp}\n" + f"Response: {output}" + ) + +print("🔡 Loading tokenizer...") +model_name = "Qwen/Qwen2-1.5B-Instruct" +tokenizer = AutoTokenizer.from_pretrained(model_name) +tokenizer.pad_token = tokenizer.eos_token + +bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + bnb_4bit_compute_dtype=torch.bfloat16, +) + +print("⚙️ Loading model with 4-bit quantization...") +model = AutoModelForCausalLM.from_pretrained( + model_name, + quantization_config=bnb_config, + device_map="auto", +) +model = prepare_model_for_kbit_training(model) + +print("🪶 Applying LoRA configuration...") +lora_config = LoraConfig( + r=64, + lora_alpha=16, + target_modules=["q_proj", "v_proj"], + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", +) +model = get_peft_model(model, lora_config) +model.print_trainable_parameters() + +print("✂️ Tokenizing dataset...") + +def tokenize_function(batch): + texts = [ + build_prompt(inst, inp, out) + for inst, inp, out in zip(batch["instruction"], batch["input"], batch["output"]) + ] + tokens = tokenizer( + texts, + padding="max_length", + truncation=True, + max_length=1024, + ) + tokens["labels"] = tokens["input_ids"].copy() + return tokens + +tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset["train"].column_names) + +args = TrainingArguments( + output_dir=save_dir, + per_device_train_batch_size=1, + gradient_accumulation_steps=4, + num_train_epochs=1, + learning_rate=2e-4, + fp16=True, + save_total_limit=2, + evaluation_strategy="epoch", + logging_dir="./logs", + logging_steps=25, + report_to="none", +) + +trainer = Trainer( + model=model, + args=args, + train_dataset=tokenized_dataset["train"], + eval_dataset=tokenized_dataset["validation"], + tokenizer=tokenizer, +) + +print("🚀 Starting training...") +trainer.train() + +print(f"💾 Saving LoRA fine-tuned model to: {save_dir}") +os.makedirs(save_dir, exist_ok=True) +model.save_pretrained(save_dir) +tokenizer.save_pretrained(save_dir) + +print(f"✅ Model and tokenizer saved successfully at: {save_dir}") \ No newline at end of file diff --git a/Processing/finetuning/thunderbird_qwen2.py b/Processing/finetuning/thunderbird_qwen2.py index a93563f..58288b0 100644 --- a/Processing/finetuning/thunderbird_qwen2.py +++ b/Processing/finetuning/thunderbird_qwen2.py @@ -1,21 +1,115 @@ -from base_finetuning import BaseFineTune -import logging +import os +import torch +from datasets import load_dataset +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + TrainingArguments, + Trainer, +) +from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training +from transformers import BitsAndBytesConfig -logger = logging.getLogger(__name__) -logger.setLevel(logging.INFO) -MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" -LOG_NAME = "thunderbird" +train_path = "/content/data/train_thunderbird.json" +test_path = "/content/data/test_thunderbird.json" -class ThunderbirdQwen2FineTune(BaseFineTune): - def __init__(self): - super().__init__(MODEL_NAME, logger=logger) +dataset_name = os.path.splitext(os.path.basename(train_path))[0].replace("train_", "") +save_dir = f"./qwen_lora_{dataset_name}_model" - def run_finetuning(self, dataset): - self.finetune(LOG_NAME, dataset) +print("🧠 Loading datasets...") +dataset = load_dataset( + "json", + data_files={"train": train_path, "validation": test_path}, +) - def save_checkpoint(self): - self.save_checkpoint(f"./{LOG_NAME}_model/checkpoint") +print(dataset) - def upload_to_hf(self): - # Implement uploading to Hugging Face logic for Thunderbird Qwen2 model - pass \ No newline at end of file +def build_prompt(instruction, inp, output): + """Format each sample into an instruction-style prompt.""" + return ( + f"Instruction: {instruction}\n" + f"Input: {inp}\n" + f"Response: {output}" + ) + +print("🔡 Loading tokenizer...") +model_name = "Qwen/Qwen2-1.5B-Instruct" +tokenizer = AutoTokenizer.from_pretrained(model_name) +tokenizer.pad_token = tokenizer.eos_token + +bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + bnb_4bit_compute_dtype=torch.bfloat16, +) + +print("⚙️ Loading model with 4-bit quantization...") +model = AutoModelForCausalLM.from_pretrained( + model_name, + quantization_config=bnb_config, + device_map="auto", +) +model = prepare_model_for_kbit_training(model) + +print("🪶 Applying LoRA configuration...") +lora_config = LoraConfig( + r=64, + lora_alpha=16, + target_modules=["q_proj", "v_proj"], + lora_dropout=0.05, + bias="none", + task_type="CAUSAL_LM", +) +model = get_peft_model(model, lora_config) +model.print_trainable_parameters() + +print("✂️ Tokenizing dataset...") + +def tokenize_function(batch): + texts = [ + build_prompt(inst, inp, out) + for inst, inp, out in zip(batch["instruction"], batch["input"], batch["output"]) + ] + tokens = tokenizer( + texts, + padding="max_length", + truncation=True, + max_length=1024, + ) + tokens["labels"] = tokens["input_ids"].copy() + return tokens + +tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset["train"].column_names) + +args = TrainingArguments( + output_dir=save_dir, + per_device_train_batch_size=1, + gradient_accumulation_steps=4, + num_train_epochs=1, + learning_rate=2e-4, + fp16=True, + save_total_limit=2, + evaluation_strategy="epoch", + logging_dir="./logs", + logging_steps=25, + report_to="none", +) + +trainer = Trainer( + model=model, + args=args, + train_dataset=tokenized_dataset["train"], + eval_dataset=tokenized_dataset["validation"], + tokenizer=tokenizer, +) + +print("🚀 Starting training...") +trainer.train() + +print(f"💾 Saving LoRA fine-tuned model to: {save_dir}") +os.makedirs(save_dir, exist_ok=True) +model.save_pretrained(save_dir) +tokenizer.save_pretrained(save_dir) + +print(f"✅ Model and tokenizer saved successfully at: {save_dir}") \ No newline at end of file diff --git a/src/SystemLogLLM/main.py b/src/SystemLogLLM/main.py index 8d1fcc0..8768993 100644 --- a/src/SystemLogLLM/main.py +++ b/src/SystemLogLLM/main.py @@ -1,4 +1,124 @@ -""" -This will be fastapi endpoint to use the SystemLogLLM model for inference. -""" +import os +import json +import torch +from tqdm import tqdm +from transformers import AutoTokenizer, AutoModelForCausalLM +from peft import PeftModel +from sklearn.metrics import accuracy_score, classification_report +BASE_MODEL = "Qwen/Qwen2-1.5B-Instruct" + +MODEL_PATHS = { + "adfa": "QWEN2/adfa", + "hdfs": "QWEN2/hdfs", + "openstack": "QWEN2/openstack", + "thunderbird": "QWEN2/thunderbird", +} + +ORCHESTRATOR_PATH = "QWEN2/orchestrator" +TEST_PATH = "data/test_adfa.json" +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +def load_lora_model(path): + """Load a Qwen base model and attach LoRA weights.""" + print(f"🔹 Loading model from {path}") + base = AutoModelForCausalLM.from_pretrained( + BASE_MODEL, + device_map="auto", + torch_dtype=torch.bfloat16, + load_in_4bit=True, + ) + model = PeftModel.from_pretrained(base, path) + model = model.merge_and_unload() + model.eval() + return model + +print("🚀 Loading models...\n") +tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) +if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + +orchestrator = load_lora_model(ORCHESTRATOR_PATH) +domain_models = {name: load_lora_model(path) for name, path in MODEL_PATHS.items()} +print(f"\n✅ Loaded {len(domain_models)} domain models and orchestrator.\n") + +with open(TEST_PATH, "r") as f: + test_data = json.load(f) +print(f"🧩 Loaded {len(test_data)} test samples.\n") + +def get_orchestrator_decision(log_sample): + """Use the orchestrator to decide which log model to route to.""" + prompt = f""" + You are a log classifier. + Decide which system this log belongs to: adfa, hdfs, openstack, thunderbird. + + Log entry: {log_sample['input']} + + Reply with only one system name. + """ + inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE) + with torch.no_grad(): + outputs = orchestrator.generate( + **inputs, + max_new_tokens=20, + temperature=0.3, + do_sample=False, + pad_token_id=tokenizer.eos_token_id, + ) + text = tokenizer.decode(outputs[0], skip_special_tokens=True).lower() + for name in MODEL_PATHS.keys(): + if name in text: + return name + return "adfa" + +def generate_response(model, log_text): + """Generate prediction (attack/normal or label) from domain model.""" + prompt = f"Analyze the following log and predict the label:\n\n{log_text}\n\nPrediction:" + inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE) + with torch.no_grad(): + outputs = model.generate( + **inputs, + max_new_tokens=20, + temperature=0.3, + pad_token_id=tokenizer.eos_token_id, + ) + result = tokenizer.decode(outputs[0], skip_special_tokens=True) + return result.split("Prediction:")[-1].strip().lower() + +predictions, gold, routes = [], [], [] + +print("🏁 Starting inference on test set...\n") + +for sample in tqdm(test_data): + try: + system = get_orchestrator_decision(sample) + model = domain_models.get(system, domain_models["adfa"]) + + pred = generate_response(model, sample["input"]) + predictions.append(pred) + gold.append(sample["output"].lower()) + routes.append(system) + except Exception as e: + print(f"⚠️ Error processing sample: {e}") + predictions.append("error") + gold.append(sample["output"].lower()) + routes.append("error") + +print("\n📊 Evaluation Summary:\n") +acc = accuracy_score(gold, predictions) +print(f"✅ Overall Accuracy: {acc:.4f}\n") + +print("🔍 Detailed Classification Report:\n") +print(classification_report(gold, predictions, zero_division=0)) + +output_results = [ + {"input": t["input"], "gold": g, "pred": p, "routed_to": r} + for t, g, p, r in zip(test_data, gold, predictions, routes) +] + +os.makedirs("results", exist_ok=True) +with open("results/inference_results.json", "w") as f: + json.dump(output_results, f, indent=2) + +print("\n💾 Saved detailed predictions to results/inference_results.json") +print("🎯 Done!\n") \ No newline at end of file