Upload folder using huggingface_hub

baseline.py

@@ -0,0 +1,90 @@
+import pandas as pd
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from sklearn.metrics import classification_report, confusion_matrix
+from tqdm import tqdm
+import os
+
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+def main():
+    device = get_device()
+    print(f"Using device: {device}")
+
+    model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
+    print(f"Loading model and tokenizer: {model_id}")
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    # Using float16 for efficiency on MPS/CUDA, or float32 on CPU
+    torch_dtype = torch.float16 if device != "cpu" else torch.float32
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch_dtype,
+        device_map=device
+    )
+
+    # Load test data
+    test_path = "data/test.csv"
+    if not os.path.exists(test_path):
+        print(f"Error: {test_path} not found. Please run data_loader.py first.")
+        return
+
+    df = pd.read_csv(test_path)
+    # To keep the baseline test fast, let's run on 100 for a quick baseline.
+    sample_size = min(100, len(df))
+    df_sample = df.sample(sample_size, random_state=42)
+
+    predictions = []
+    labels = []
+
+    print(f"Evaluating zero-shot performance on {sample_size} samples...")
+
+    for _, row in tqdm(df_sample.iterrows(), total=sample_size):
+        text = str(row['text'])
+        label = int(row['phishing']) # 0 for safe, 1 for phishing
+
+        # SmolLM2-Instruct prompt format
+        messages = [{"role": "user", "content": f"""Classify the following email text as either 'Safe' or 'Phishing'. Respond with only one word: 'Safe' or 'Phishing'.
+
+Email text: {text}
+
+Classification:"""}]
+        input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        
+        inputs = tokenizer(input_text, return_tensors="pt").to(device)
+        
+        with torch.no_grad():
+            output = model.generate(
+                **inputs, 
+                max_new_tokens=10, 
+                temperature=0.1, 
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        response = tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip().lower()
+        
+        if 'phishing' in response:
+            predictions.append(1)
+        elif 'safe' in response:
+            predictions.append(0)
+        else:
+            # Fallback if the model doesn't follow instructions well
+            # print(f"Warning: Model gave unexpected response: '{response}'")
+            predictions.append(0) 
+
+        labels.append(label)
+
+    print("\nBaseline Results (Zero-Shot):")
+    print(classification_report(labels, predictions, target_names=['Safe', 'Phishing'], zero_division=0))
+    print("\nConfusion Matrix:")
+    print(confusion_matrix(labels, predictions))
+
+if __name__ == "__main__":
+    main()

data_loader.py

@@ -0,0 +1,43 @@
+import os
+from datasets import load_dataset
+from sklearn.model_selection import train_test_split
+import pandas as pd
+
+
+def prepare_data(dataset_name: str = "David-Egea/phishing-texts"):
+    print(f"Loading dataset: {dataset_name}...")
+    # The dataset usually loads into a 'train' split if not specified
+    ds = load_dataset(dataset_name)
+
+    # Convert to pandas for easier manipulation/splitting
+    df: pd.DataFrame = ds["train"].to_pandas()  # type: ignore
+
+    print(f"Total samples: {len(df)}")
+    print(f"Class distribution:\n{df['phishing'].value_counts(normalize=True)}")
+
+    # 80% Train, 20% Temp (Val + Test)
+    train_df, temp_df = train_test_split(
+        df, test_size=0.2, random_state=42, stratify=df["phishing"]
+    )
+
+    # Split temp into 50% Val, 50% Test (results in 10% each of total)
+    val_df, test_df = train_test_split(
+        temp_df, test_size=0.5, random_state=42, stratify=temp_df["phishing"]
+    )
+
+    print(f"Train samples: {len(train_df)}")
+    print(f"Val samples: {len(val_df)}")
+    print(f"Test samples: {len(test_df)}")
+
+    # Ensure data directory exists
+    os.makedirs("data", exist_ok=True)
+
+    # Save splits
+    train_df.to_csv("data/train.csv", index=False)
+    val_df.to_csv("data/val.csv", index=False)
+    test_df.to_csv("data/test.csv", index=False)
+    print("Splits saved to data/ folder.")
+
+
+if __name__ == "__main__":
+    prepare_data()

train.py

@@ -0,0 +1,164 @@
+import os
+import torch
+from datasets import load_dataset, Dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+)
+from peft import LoraConfig
+from trl.trainer.sft_trainer import SFTTrainer
+from trl.trainer.sft_config import SFTConfig
+import argparse
+import pandas as pd
+
+# Define tokenizer globally for the mapping function
+tokenizer = None
+
+
+def format_instruction(sample):
+    # Standard format for SmolLM2-Instruct
+    label_str = "Phishing" if sample["phishing"] == 1 else "Safe"
+
+    messages = [
+        {
+            "role": "user",
+            "content": f"Classify the following email text as either 'Safe' or 'Phishing'. Respond with only one word: 'Safe' or 'Phishing'.\n\nEmail text: {sample['text']}\n\nClassification:",
+        },
+        {"role": "assistant", "content": label_str},
+    ]
+    # tokenizer is now accessible globally
+    return (
+        {"text": tokenizer.apply_chat_template(messages, tokenize=False)}
+        if tokenizer
+        else {"text": ""}
+    )
+
+
+def main(args):
+    global tokenizer
+    device = (
+        "cuda"
+        if torch.cuda.is_available()
+        else "mps"
+        if torch.backends.mps.is_available()
+        else "cpu"
+    )
+    print(f"Using device: {device}")
+
+    model_id = args.model_id
+    print(f"Loading tokenizer and model: {model_id}")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    # Load Model
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        device_map=device if device != "mps" else None,
+    )
+    if device == "mps":
+        model.to("mps")  # type: ignore
+
+    # LoRA Configuration
+    peft_config = LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        target_modules=[
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+
+    # Load Data
+    print(f"Loading data from {args.dataset_name}...")
+    if os.path.exists(args.dataset_name):
+        train_df = pd.read_csv(os.path.join(args.dataset_name, "train.csv"))
+        val_df = pd.read_csv(os.path.join(args.dataset_name, "val.csv"))
+        if args.quick_test:
+            train_df = train_df.head(100)
+            val_df = val_df.head(20)
+        train_dataset = Dataset.from_pandas(train_df)
+        val_dataset = Dataset.from_pandas(val_df)
+    else:
+        dataset = load_dataset(args.dataset_name)
+        train_dataset = dataset["train"]
+        val_dataset = dataset["validation"] if "validation" in dataset else None
+
+    # Apply formatting
+    print("Formatting datasets...")
+    train_dataset = train_dataset.map(format_instruction)
+    if val_dataset:
+        val_dataset = val_dataset.map(format_instruction)
+
+    # Use SFTConfig for modern TRL
+    sft_config = SFTConfig(
+        output_dir=args.output_dir,
+        per_device_train_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.grad_accum,
+        learning_rate=args.lr,
+        logging_steps=10,
+        num_train_epochs=args.epochs,
+        max_steps=args.max_steps,
+        eval_strategy="steps" if val_dataset else "no",
+        eval_steps=100,
+        save_strategy="steps",
+        save_steps=100,
+        lr_scheduler_type="cosine",
+        warmup_ratio=0.1,
+        bf16=torch.cuda.is_available(),
+        push_to_hub=args.push_to_hub,
+        report_to="tensorboard" if not args.no_report else "none",
+        remove_unused_columns=False,
+        dataset_text_field="text",
+        max_length=args.max_seq_length,
+    )
+
+    # Standard HF SFTTrainer
+    trainer = SFTTrainer(
+        model=model,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        peft_config=peft_config,
+        processing_class=tokenizer,
+        args=sft_config,
+    )
+
+    print("Starting training...")
+    trainer.train()
+
+    print(f"Saving model to {args.output_dir}")
+    trainer.save_model(args.output_dir)
+    if args.push_to_hub:
+        trainer.push_to_hub()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_id", type=str, default="HuggingFaceTB/SmolLM2-135M-Instruct"
+    )
+    parser.add_argument("--dataset_name", type=str, default="data/")
+    parser.add_argument("--output_dir", type=str, default="models/smollm2-phish-sft")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--grad_accum", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=2e-4)
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--max_steps", type=int, default=-1)
+    parser.add_argument("--max_seq_length", type=int, default=512)
+    parser.add_argument("--lora_r", type=int, default=16)
+    parser.add_argument("--lora_alpha", type=int, default=32)
+    parser.add_argument("--lora_dropout", type=float, default=0.05)
+    parser.add_argument("--quick_test", action="store_true")
+    parser.add_argument("--push_to_hub", action="store_true")
+    parser.add_argument("--no_report", action="store_true")
+    args = parser.parse_args()
+    main(args)

verify_setup.py

@@ -0,0 +1,29 @@
+import torch
+import transformers
+import datasets
+import peft
+import accelerate
+import platform
+
+
+def verify():
+    print(f"OS: {platform.system()} {platform.release()}")
+    print(f"Python: {platform.python_version()}")
+    print("-" * 20)
+    print(f"PyTorch version: {torch.__version__}")
+    print(f"Transformers version: {transformers.__version__}")
+    print(f"Datasets version: {datasets.__version__}")
+    print(f"PEFT version: {peft.__version__}")
+    print(f"Accelerate version: {accelerate.__version__}")
+
+    # Check for GPU
+    if torch.cuda.is_available():
+        print(f"GPU: {torch.cuda.get_device_name(0)} (CUDA available)")
+    elif torch.backends.mps.is_available():
+        print("GPU: Apple Silicon MPS available")
+    else:
+        print("GPU: Not available (using CPU)")
+
+
+if __name__ == "__main__":
+    verify()