Create README.md

README.md

@@ -0,0 +1,141 @@
+---
+license: apache-2.0
+language:
+- en
+base_model:
+- facebook/opt-1.3b
+tags:
+- text-generation
+- peft
+- aws
+- lora
+- security
+- iam
+- fine-tuned
+---
+Model Card: secure-policy-rewriter 🔐
+This is a fine-tuned version of the facebook/opt-1.3b model designed to analyze and rewrite risky AWS IAM policies into more secure, specific policies based on the Principle of Least Privilege.
+
+The model takes a potentially risky IAM policy (often with wildcards like "*" or "s3:*") as input and generates a secure version with specific permissions, resources, and optional conditions. This model can be used as a component in a larger security scanning or remediation pipeline.
+
+Model Details
+Base Model: facebook/opt-1.3b
+
+Task: Text Generation (text-generation)
+
+Fine-tuning Method: Parameter-Efficient Fine-Tuning (PEFT) using LoRA (Low-Rank Adaptation)
+
+Libraries: The model was trained using transformers, datasets, peft, and accelerate
+
+Training Data
+The model was fine-tuned on a custom dataset of 100000 JSON pairs, with the input being a risky policy and the output being a securely rewritten version. The dataset was generated programmatically within a Kaggle Notebook and contained a mix of policies with specific actions and wildcards for both actions and resources.
+
+How to Use 💻
+To use this model, you must first load the base model, then attach the PEFT adapter weights.
+
+Dependencies:
+Install the necessary libraries:
+
+Bash
+
+pip install transformers torch peft accelerate
+Inference Code:
+The following Python script demonstrates how to load the model and perform a rewrite on a sample policy.
+
+Python
+
+import json
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from peft import PeftModel, LoraConfig, TaskType
+
+# Set device
+device = 0 if torch.cuda.is_available() else -1
+
+# Load base model and tokenizer
+base_model_id = "facebook/opt-1.3b"
+checkpoint_path = "Yehia3A/secure-policy-rewriter" 
+
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
+model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
+
+# Load LoRA adapter
+lora_config = LoraConfig(
+    task_type=TaskType.CAUSAL_LM,
+    r=16,
+    lora_alpha=32,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    inference_mode=True
+)
+model = PeftModel.from_pretrained(model, checkpoint_path, config=lora_config)
+model.eval()
+
+# Create a text generation pipeline
+llm = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=device,
+    pad_token_id=tokenizer.eos_token_id
+)
+
+def rewrite_policy(policy_json: dict) -> str:
+    # Use a few-shot prompt to guide the model
+    prompt = f"""Rewrite risky IAM policies to be secure. Replace wildcards with specific permissions and add conditions.
+
+Input: {json.dumps(policy_json, indent=2)}
+
+Output:"""
+    
+    result = llm(prompt, max_new_tokens=300, temperature=0.05, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+    result_text = result[0]["generated_text"].strip()
+    
+    # Simple JSON extraction from the output
+    try:
+        start_index = result_text.find("{")
+        end_index = result_text.rfind("}") + 1
+        if start_index != -1 and end_index != -1:
+            json_str = result_text[start_index:end_index]
+            return json.dumps(json.loads(json_str), indent=2)
+    except:
+        return "{}"
+    return "{}"
+
+
+# Example risky policy
+risky_policy = {
+  "Version": "2012-10-17",
+  "Statement": [{
+    "Effect": "Allow",
+    "Action": "*",
+    "Resource": "*"
+  }]
+}
+
+# Rewrite the policy
+rewritten_policy = rewrite_policy(risky_policy)
+print("Rewritten Secure Policy:")
+print(rewritten_policy)
+
+Files
+This repository contains the necessary files to run the fine-tuned model:
+
+adapter_config.json: Configuration for the PEFT adapter.
+
+adapter_model.safetensors: The fine-tuned weights of the PEFT adapter.
+
+tokenizer_config.json: The tokenizer configuration.
+
+tokenizer.json: The tokenizer model file.
+
+special_tokens_map.json: A mapping of special tokens.
+
+merges.txt: A vocabulary file for the tokenizer.
+
+Disclaimers
+This model is a proof-of-concept for educational and research purposes. It is not intended for use in production environments without further validation, testing, and security checks. It was trained on a small, synthetic dataset and may not generalize well to all real-world scenarios.