Add inference script and remove draft status heading

README.md

@@ -37,11 +37,6 @@ separate external source.
 This repository contains a LoRA adapter, not the full IBM base model.
 This is an independent student project, not an IBM-endorsed detector.
 
-## Release status
-
-Private draft. Fresh-adapter inference verification and standalone usage
-instructions are pending. This section will be updated before public release.
-
 ## Model and training
 
 | Setting | Value |

inference.py

@@ -0,0 +1,197 @@
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from peft import PeftModel, prepare_model_for_kbit_training
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+
+class URLClassifier:
+    def __init__(self, repository, revision):
+        if not torch.cuda.is_available():
+            raise RuntimeError("A CUDA GPU is required.")
+        if not torch.cuda.is_bf16_supported():
+            raise RuntimeError("A BF16-capable GPU is required.")
+
+        root = Path(snapshot_download(
+            repo_id=repository,
+            revision=revision,
+            allow_patterns=[
+                "adapter_config.json",
+                "adapter_model.safetensors",
+                "tokenizer.json",
+                "tokenizer_config.json",
+                "special_tokens_map.json",
+                "added_tokens.json",
+                "chat_template.jinja",
+                "experiment_config/*.json",
+            ],
+        ))
+
+        def read_config(name):
+            return json.loads(
+                (root / "experiment_config" / name).read_text(
+                    encoding="utf-8"
+                )
+            )
+
+        base = read_config("base_model.json")
+        self.policy = read_config("input_policy.json")
+        prompt = read_config("prompt_config.json")
+
+        self.instruction = prompt["instruction"]
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            str(root), local_files_only=True
+        )
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "left"
+
+        self.label_ids = prompt["label_token_ids"]
+        for label in ["A", "B"]:
+            ids = self.tokenizer.encode(label, add_special_tokens=False)
+            if ids != [self.label_ids[label]]:
+                raise RuntimeError(f"Unexpected tokenization for {label}.")
+
+        quantization = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+
+        model = AutoModelForCausalLM.from_pretrained(
+            base["model_id"],
+            revision=base["revision"],
+            quantization_config=quantization,
+            device_map={"": 0},
+            dtype=torch.bfloat16,
+            attn_implementation="sdpa",
+        )
+        model = prepare_model_for_kbit_training(
+            model, use_gradient_checkpointing=False
+        )
+        self.model = PeftModel.from_pretrained(
+            model, str(root), is_trainable=False
+        )
+        self.model.eval()
+        self.model.config.use_cache = False
+
+    def _encode(self, url):
+        if not isinstance(url, str) or not url.strip():
+            raise ValueError("URL must be a nonempty string.")
+
+        def encode(text):
+            rendered = self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": self.instruction + text}],
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False,
+            )
+            return self.tokenizer.encode(
+                rendered, add_special_tokens=False
+            )
+
+        prompt_ids = encode(url)
+        original_length = len(prompt_ids)
+        maximum = self.policy["max_prompt_tokens"]
+
+        if original_length <= maximum:
+            return prompt_ids, False
+
+        url_ids = self.tokenizer.encode(url, add_special_tokens=False)
+
+        while len(prompt_ids) > maximum:
+            excess = len(prompt_ids) - maximum
+            keep = max(0, len(url_ids) - excess - 8)
+
+            if keep >= len(url_ids):
+                raise RuntimeError("Truncation made no progress.")
+
+            url_ids = url_ids[:keep]
+            text = self.tokenizer.decode(
+                url_ids,
+                skip_special_tokens=False,
+                clean_up_tokenization_spaces=False,
+            )
+            prompt_ids = encode(text)
+
+            if not url_ids and len(prompt_ids) > maximum:
+                raise ValueError("Instructions exceed the token budget.")
+
+        return prompt_ids, True
+
+    def predict_urls(self, urls, batch_size=4):
+        if not isinstance(batch_size, int) or batch_size < 1:
+            raise ValueError("batch_size must be a positive integer.")
+
+        urls = list(urls)
+        results = []
+
+        for start in range(0, len(urls), batch_size):
+            encoded = [
+                self._encode(url)
+                for url in urls[start:start + batch_size]
+            ]
+            features = [
+                {
+                    "input_ids": ids,
+                    "attention_mask": [1] * len(ids),
+                }
+                for ids, truncated in encoded
+            ]
+
+            inputs = self.tokenizer.pad(
+                features, padding=True, return_tensors="pt"
+            ).to("cuda")
+
+            with torch.inference_mode():
+                output = self.model(**inputs, use_cache=False)
+                logits = output.logits[
+                    :, -1, [self.label_ids["A"], self.label_ids["B"]]
+                ].float()
+
+                if not torch.isfinite(logits).all().item():
+                    raise RuntimeError("Non-finite label logits.")
+
+                margins = logits[:, 1] - logits[:, 0]
+                scores = torch.softmax(logits, dim=-1)[:, 1]
+
+            for margin, score, (_, truncated) in zip(
+                margins.cpu().tolist(),
+                scores.cpu().tolist(),
+                encoded,
+            ):
+                prediction = 0 if margin > 0 else 1
+                results.append({
+                    "label": "phishing" if prediction == 0 else "legitimate",
+                    "prediction": prediction,
+                    "phishing_score_uncalibrated": score,
+                    "phishing_logit_margin": margin,
+                    "url_truncated": truncated,
+                })
+
+            del inputs, output, logits, margins, scores
+
+        return results
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--repository",
+        default="bengoldberg0/granite-4.2-3b-phishing-url-qlora",
+    )
+    parser.add_argument("--revision", required=True)
+    parser.add_argument("--url", required=True)
+    args = parser.parse_args()
+
+    classifier = URLClassifier(args.repository, args.revision)
+    result = classifier.predict_urls([args.url])[0]
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,8 @@
+torch==2.11.0+cu128
+transformers==5.17.0
+peft==0.21.0
+accelerate==1.15.0
+bitsandbytes==0.50.2
+huggingface_hub==1.31.0
+safetensors==0.8.0
+tokenizers==0.23.1