Create open_ended_question_generator_secure.py

open_ended_question_generator_secure.py

@@ -0,0 +1,410 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+open_ended_question_generator_secure.py
+
+End-to-end script to generate open-ended questions from context(s) with:
+- Robust list-formatted parsing
+- CLI with single or batch inputs (TXT/CSV)
+- Reproducibility (seed)
+- Device auto-select (CUDA / MPS / CPU)
+- Export to JSON / CSV / TXT
+- Optional AES-256-like authenticated encryption via Fernet (with PBKDF2 key derivation)
+- Optional decryption utility
+
+Dependencies:
+  pip install torch transformers cryptography
+
+Example:
+  python open_ended_question_generator_secure.py \
+    --context "AGI for cosmology" --n 5 --model gpt2-large \
+    --out questions.json --format json --encrypt --password "your-secret"
+"""
+
+import os
+import re
+import csv
+import json
+import argparse
+import getpass
+import base64
+import sys
+from typing import List, Dict, Tuple, Optional
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# --- Optional encryption deps ---
+try:
+    from cryptography.fernet import Fernet
+    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+    from cryptography.hazmat.primitives import hashes
+    from cryptography.hazmat.backends import default_backend
+except Exception:
+    Fernet = None  # Will validate at runtime if encryption/decryption is used.
+
+
+# ----------------------------
+# Device selection
+# ----------------------------
+def select_device() -> torch.device:
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+# ----------------------------
+# Prompt and parsing
+# ----------------------------
+PROMPT_TEMPLATE = """You are a master at generating deep, open-ended, and thought-provoking questions.
+Each question must be:
+- Self-contained and understandable without extra context.
+- Exploratory (not answerable with yes/no).
+- Written in clear, engaging language.
+
+Context:
+{context}
+
+Output exactly {n} questions as a numbered list, one per line, formatted like:
+1. ...
+2. ...
+3. ...
+No extra commentary, no headings, no explanations — just the list.
+"""
+
+def build_prompt(context: str, n: int) -> str:
+    return PROMPT_TEMPLATE.format(context=context.strip(), n=n)
+
+_Q_LINE_RE = re.compile(r"^\s*(\d+)\.\s+(.*\S)\s*$")
+
+def normalize_q(q: str) -> str:
+    q = q.strip()
+    # Ensure it ends with a question mark for consistency
+    if not q.endswith("?"):
+        q += "?"
+    return q
+
+def parse_questions_from_text(text: str, n: int) -> List[str]:
+    lines = text.splitlines()
+    candidates = []
+    for line in lines:
+        m = _Q_LINE_RE.match(line)
+        if m:
+            q_text = normalize_q(m.group(2))
+            candidates.append(q_text)
+    # Deduplicate while preserving order
+    seen = set()
+    unique = []
+    for q in candidates:
+        key = q.lower().strip()
+        if key not in seen:
+            seen.add(key)
+            unique.append(q)
+    return unique[:n]
+
+
+# ----------------------------
+# Model loading and generation
+# ----------------------------
+def load_model_and_tokenizer(model_name: str, device: torch.device):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model.to(device)
+    # For models like GPT-2 without a pad token
+    if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    return model, tokenizer
+
+def generate_questions_once(
+    model,
+    tokenizer,
+    device: torch.device,
+    context: str,
+    n: int,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+) -> List[str]:
+    prompt = build_prompt(context, n)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    output = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        do_sample=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    # Extract only the continuation after the prompt
+    # In many causal LMs, decoded contains prompt + completion; we slice from len(input_ids)
+    # Simpler approach: parse all lines and trust the numbered format.
+    questions = parse_questions_from_text(decoded, n)
+    return questions
+
+def generate_questions(
+    model,
+    tokenizer,
+    device: torch.device,
+    context: str,
+    n: int = 3,
+    max_new_tokens: int = 200,
+    temperature: float = 0.95,
+    top_p: float = 0.95,
+    top_k: int = 50,
+    seed: Optional[int] = None,
+    attempts: int = 3,
+) -> List[str]:
+    if seed is not None:
+        torch.manual_seed(seed)
+        if device.type == "cuda":
+            torch.cuda.manual_seed_all(seed)
+    collected: List[str] = []
+    tried = 0
+    while len(collected) < n and tried < attempts:
+        tried += 1
+        # Slightly adjust temperature on retries to improve variety
+        temp = min(1.2, max(0.7, temperature + 0.1 * (tried - 1)))
+        qs = generate_questions_once(
+            model, tokenizer, device, context, n, max_new_tokens, temp, top_p, top_k
+        )
+        # Merge unique
+        existing = set([q.lower().strip() for q in collected])
+        for q in qs:
+            key = q.lower().strip()
+            if key not in existing and len(collected) < n:
+                collected.append(q)
+                existing.add(key)
+    # If still short, pad with simple variants (rare)
+    while len(collected) < n:
+        collected.append(collected[-1] + " (expand)") if collected else collected.append("What deeper questions arise from this context?")
+    return collected[:n]
+
+
+# ----------------------------
+# Batch input handling
+# ----------------------------
+def load_contexts(source_text: Optional[str], source_file: Optional[str]) -> List[Tuple[str, str]]:
+    """
+    Returns list of (context_id, context_text).
+    - If source_text is provided, returns single-item list.
+    - If CSV file: expects a 'context' column.
+    - If TXT/MD: splits on lines containing only '---' or returns whole file as one context.
+    """
+    out: List[Tuple[str, str]] = []
+    if source_text:
+        out.append(("context_1", source_text.strip()))
+        return out
+    if not source_file:
+        raise ValueError("Either --context or --context-file is required.")
+    if not os.path.exists(source_file):
+        raise FileNotFoundError(f"Context file not found: {source_file}")
+
+    ext = os.path.splitext(source_file)[1].lower()
+    if ext == ".csv":
+        with open(source_file, "r", encoding="utf-8", newline="") as f:
+            reader = csv.DictReader(f)
+            if "context" not in reader.fieldnames:
+                raise ValueError("CSV must have a 'context' column.")
+            for i, row in enumerate(reader, start=1):
+                ctx = (row.get("context") or "").strip()
+                if ctx:
+                    out.append((f"context_{i}", ctx))
+    else:
+        # Plain text / markdown: split on '---' delimiter lines if present
+        with open(source_file, "r", encoding="utf-8") as f:
+            content = f.read()
+        parts = re.split(r"^\s*---\s*$", content, flags=re.MULTILINE)
+        parts = [p.strip() for p in parts if p.strip()]
+        if not parts:
+            raise ValueError("No context found in file.")
+        for i, ctx in enumerate(parts, start=1):
+            out.append((f"context_{i}", ctx))
+    return out
+
+
+# ----------------------------
+# Output writers
+# ----------------------------
+def write_json(out_path: str, rows: List[Dict]):
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump(rows, f, ensure_ascii=False, indent=2)
+
+def write_csv(out_path: str, rows: List[Dict], n: int):
+    fieldnames = ["context_id", "context"] + [f"q{i}" for i in range(1, n + 1)]
+    with open(out_path, "w", encoding="utf-8", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for r in rows:
+            writer.writerow(r)
+
+def write_txt(out_path: str, rows: List[Dict], n: int):
+    with open(out_path, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(f"[{r['context_id']}]\n")
+            f.write(r["context"].strip() + "\n")
+            for i in range(1, n + 1):
+                f.write(f"{i}. {r[f'q{i}']}\n")
+            f.write("\n")
+
+
+# ----------------------------
+# Encryption / Decryption
+# ----------------------------
+MAGIC = b"QSEC1"
+
+def require_crypto():
+    if Fernet is None:
+        raise RuntimeError("Encryption requested but 'cryptography' is not installed. Run: pip install cryptography")
+
+def derive_key_from_password(password: str, salt: bytes) -> bytes:
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        iterations=200_000,
+        backend=default_backend(),
+    )
+    key = kdf.derive(password.encode("utf-8"))
+    return base64.urlsafe_b64encode(key)
+
+def encrypt_file(in_path: str, out_path: str, password: str):
+    require_crypto()
+    with open(in_path, "rb") as f:
+        plaintext = f.read()
+    salt = os.urandom(16)
+    key = derive_key_from_password(password, salt)
+    fernet = Fernet(key)
+    ciphertext = fernet.encrypt(plaintext)
+    with open(out_path, "wb") as f:
+        f.write(MAGIC + salt + ciphertext)
+
+def decrypt_file(in_path: str, out_path: str, password: str):
+    require_crypto()
+    with open(in_path, "rb") as f:
+        blob = f.read()
+    if not blob.startswith(MAGIC) or len(blob) < len(MAGIC) + 16 + 1:
+        raise ValueError("Invalid or unsupported encrypted file.")
+    salt = blob[len(MAGIC):len(MAGIC)+16]
+    ciphertext = blob[len(MAGIC)+16:]
+    key = derive_key_from_password(password, salt)
+    fernet = Fernet(key)
+    plaintext = fernet.decrypt(ciphertext)
+    with open(out_path, "wb") as f:
+        f.write(plaintext)
+
+
+# ----------------------------
+# Main CLI
+# ----------------------------
+def main():
+    parser = argparse.ArgumentParser(description="Generate deep open-ended questions with optional encryption/decryption.")
+    mode = parser.add_mutuallyExclusiveGroup(required=True)
+    mode.add_argument("--generate", action="store_true", help="Generate questions from context(s).")
+    mode.add_argument("--decrypt", action="store_true", help="Decrypt an encrypted file (no generation).")
+
+    # Generation inputs
+    parser.add_argument("--context", type=str, help="Inline context text.")
+    parser.add_argument("--context-file", type=str, help="Path to TXT/MD (split by ---) or CSV with 'context' column.")
+    parser.add_argument("--n", type=int, default=3, help="Number of questions to generate per context.")
+    parser.add_argument("--model", type=str, default="gpt2-large", help="HuggingFace model name.")
+    parser.add_argument("--max-new-tokens", type=int, default=220, help="Max new tokens for generation.")
+    parser.add_argument("--temperature", type=float, default=0.95, help="Sampling temperature.")
+    parser.add_argument("--top-p", type=float, default=0.95, help="Top-p nucleus sampling.")
+    parser.add_argument("--top-k", type=int, default=50, help="Top-k sampling.")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
+    parser.add_argument("--attempts", type=int, default=3, help="Max attempts to reach exactly n questions.")
+
+    # Output
+    parser.add_argument("--out", type=str, default=None, help="Output file path. If omitted, prints to stdout.")
+    parser.add_argument("--format", type=str, choices=["json", "csv", "txt"], default="json", help="Output format when generating.")
+    parser.add_argument("--encrypt", action="store_true", help="Encrypt the output file after generation.")
+    parser.add_argument("--password", type=str, default=None, help="Password for encryption/decryption. If omitted, prompts securely.")
+
+    # Decryption I/O
+    parser.add_argument("--in", dest="in_path", type=str, help="Input file for decryption (encrypted).")
+    parser.add_argument("--out-decrypted", type=str, help="Output file for decrypted plaintext.")
+
+    args = parser.parse_args()
+
+    device = select_device()
+
+    if args.decrypt:
+        # Decrypt mode
+        if not args.in_path or not args.out_decrypted:
+            parser.error("--decrypt requires --in and --out-decrypted.")
+        password = args.password or getpass.getpass("Enter password: ")
+        decrypt_file(args.in_path, args.out_decrypted, password)
+        print(f"Decrypted to: {args.out_decrypted}")
+        return
+
+    # Generate mode
+    contexts = load_contexts(args.context, args.context_file)
+    model, tokenizer = load_model_and_tokenizer(args.model, device)
+
+    rows: List[Dict] = []
+    for ctx_id, ctx in contexts:
+        qs = generate_questions(
+            model=model,
+            tokenizer=tokenizer,
+            device=device,
+            context=ctx,
+            n=args.n,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            top_k=args.top_k,
+            seed=args.seed,
+            attempts=args.attempts,
+        )
+        row = {"context_id": ctx_id, "context": ctx}
+        for i, q in enumerate(qs, start=1):
+            row[f"q{i}"] = q
+        rows.append(row)
+
+    # Output
+    if args.out:
+        out_path = args.out
+        os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+        if args.format == "json":
+            write_json(out_path, rows)
+        elif args.format == "csv":
+            write_csv(out_path, rows, args.n)
+        else:
+            write_txt(out_path, rows, args.n)
+
+        if args.encrypt:
+            password = args.password or getpass.getpass("Enter password: ")
+            enc_path = out_path + ".enc"
+            encrypt_file(out_path, enc_path, password)
+            print(f"Saved: {out_path}")
+            print(f"Encrypted copy: {enc_path}")
+        else:
+            print(f"Saved: {out_path}")
+    else:
+        # Print to stdout in selected format
+        if args.format == "json":
+            print(json.dumps(rows, ensure_ascii=False, indent=2))
+        elif args.format == "csv":
+            # Minimal CSV to stdout
+            fieldnames = ["context_id", "context"] + [f"q{i}" for i in range(1, args.n + 1)]
+            writer = csv.DictWriter(sys.stdout, fieldnames=fieldnames)
+            writer.writeheader()
+            for r in rows:
+                writer.writerow(r)
+        else:
+            for r in rows:
+                print(f"[{r['context_id']}]")
+                print(r["context"].strip())
+                for i in range(1, args.n + 1):
+                    print(f"{i}. {r[f'q{i}']}")
+                print()
+
+if __name__ == "__main__":
+    main()