Spaces:
Running on Zero
Running on Zero
| import gradio as gr | |
| import torch, json, os, re | |
| import spaces | |
| from pathlib import Path | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| MODEL_IDS = { | |
| "PML-12L-O": "K0D3IN/PML-12L-O", | |
| "PML-22L-O": "K0D3IN/PML-22L-O", | |
| } | |
| COUNTRY_NAMES = [ | |
| "Turkey", "United States", "United Kingdom", "Germany", "France", "Italy", | |
| "Spain", "Netherlands", "Belgium", "Switzerland", "Austria", "Sweden", | |
| "Norway", "Denmark", "Finland", "Poland", "Czech Republic", "Hungary", | |
| "Romania", "Bulgaria", "Greece", "Croatia", "Serbia", | |
| "Brazil", "Argentina", "Chile", "Colombia", "Mexico", "Peru", | |
| "India", "China", "Japan", "South Korea", "Singapore", "Indonesia", | |
| "Philippines", "Thailand", "Vietnam", "Russia", "Ukraine", | |
| "Australia", "New Zealand", "South Africa", "Egypt", "Morocco", "Nigeria", | |
| "Israel", "Saudi Arabia", "UAE", "Canada", "Ireland", "Portugal", | |
| ] | |
| COUNTRY_CODE = {n: c for n, c in zip(COUNTRY_NAMES, [ | |
| "TR","US","GB","DE","FR","IT","ES","NL","BE","CH","AT","SE","NO","DK","FI", | |
| "PL","CZ","HU","RO","BG","GR","HR","RS","BR","AR","CL","CO","MX","PE", | |
| "IN","CN","JP","KR","SG","ID","PH","TH","VN","RU","UA","AU","NZ","ZA", | |
| "EG","MA","NG","IL","SA","AE","CA","IE","PT", | |
| ])} | |
| MODELS = {} | |
| def load_model(name): | |
| if name not in MODELS: | |
| from model_v5 import PasswordLLaMA | |
| from tokenizers import Tokenizer | |
| import safetensors.torch | |
| device = "cpu" | |
| repo_id = MODEL_IDS[name] | |
| from huggingface_hub import hf_hub_download | |
| config_path = hf_hub_download(repo_id=repo_id, filename="config.json", token=HF_TOKEN) | |
| with open(config_path) as f: | |
| cfg = json.load(f) | |
| model = PasswordLLaMA( | |
| vocab_size=cfg.get("vocab_size", 8192), | |
| n_layer=cfg.get("n_layer", 22), | |
| n_embd=cfg.get("n_embd", 384), | |
| n_head=cfg.get("n_head", 6), | |
| max_seq_len=cfg.get("max_seq_len", 48), | |
| ) | |
| weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", token=HF_TOKEN) | |
| state = safetensors.torch.load_file(weights_path, device=device) | |
| model.load_state_dict(state, strict=True) | |
| model.eval() | |
| toker_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.json", token=HF_TOKEN) | |
| toker = Tokenizer.from_file(toker_path) | |
| MODELS[name] = (model, toker) | |
| return MODELS[name] | |
| def generate(country, username, pw_len, model_name, num_pw): | |
| if not username or not username.strip(): | |
| return "Please enter a username" | |
| cc = COUNTRY_CODE.get(country, "US") | |
| username = username.strip() | |
| pw_len = int(pw_len) if pw_len and pw_len > 0 else 0 | |
| model, toker = load_model(model_name) | |
| if pw_len > 0: | |
| prompt = f"[USER:{username}][COUNTRY:{cc}][LEN:{pw_len}]:" | |
| else: | |
| prompt = f"[USER:{username}][COUNTRY:{cc}]:" | |
| prefix_ids = toker.encode(prompt).ids | |
| device = next(model.parameters()).device | |
| pad_id = toker.token_to_id("<PAD>") | |
| eos_id = toker.token_to_id("<EOS>") | |
| results = [] | |
| seen = set() | |
| bs = min(num_pw, 64) | |
| with torch.no_grad(): | |
| while len(results) < num_pw: | |
| ids = torch.full((bs, 48), pad_id, dtype=torch.long, device=device) | |
| ids[:, :len(prefix_ids)] = torch.tensor(prefix_ids, device=device) | |
| cur_len = len(prefix_ids) | |
| finished = torch.zeros(bs, dtype=torch.bool, device=device) | |
| max_new = 16 if pw_len == 0 else pw_len + 4 | |
| for _ in range(max_new): | |
| if finished.all(): | |
| break | |
| logits = model(ids[:, :cur_len]) | |
| nxt = logits[torch.arange(bs), -1, :] / 0.8 | |
| vals, _ = torch.topk(nxt, 50) | |
| nxt[nxt < vals[:, -1:]] = float("-inf") | |
| probs = torch.softmax(nxt, dim=-1) | |
| nids = torch.multinomial(probs, 1).squeeze(-1) | |
| finished |= (nids == eos_id) | |
| nids[finished] = pad_id | |
| ids[:, cur_len] = nids | |
| cur_len += 1 | |
| if pw_len > 0 and cur_len - len(prefix_ids) >= pw_len: | |
| break | |
| for i in range(bs): | |
| pw = toker.decode(ids[i].tolist()) | |
| for t in ["<BOS>", "<EOS>", "<PAD>", "<UNK>"]: | |
| pw = pw.replace(t, "") | |
| pw = re.sub(r"\[[A-Z]+:[^\]]*\]", "", pw) | |
| pw = re.sub(r"^\s*:\s*", "", pw) # artık tag ayracı | |
| pw = pw.strip() | |
| if pw and pw not in seen and len(pw) >= (pw_len if pw_len > 0 else 4): | |
| seen.add(pw) | |
| results.append(pw) | |
| return "\n".join(results[:num_pw]) | |
| with gr.Blocks(title="PML Password Generator") as demo: | |
| gr.Markdown("# 🔐 PML Password Generator\nGenerate culturally-aware passwords with USER+COUNTRY conditioning.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| country = gr.Dropdown(choices=COUNTRY_NAMES, value="Turkey", label="Country") | |
| username = gr.Textbox(label="Username", placeholder="john_doe") | |
| length = gr.Slider(0, 30, 0, step=1, label="Length (0 = auto)") | |
| model_sel = gr.Dropdown(choices=["PML-12L-O", "PML-22L-O"], value="PML-22L-O", label="Model") | |
| count = gr.Slider(1, 50, 10, step=1, label="How many?") | |
| btn = gr.Button("🎲 Generate", variant="primary") | |
| with gr.Column(scale=2): | |
| out = gr.Textbox(lines=20, label="Passwords", placeholder="Generated passwords...") | |
| btn.click(fn=generate, inputs=[country, username, length, model_sel, count], outputs=out) | |
| gr.Markdown("---\n### Models\n[PML-12L-O](https://huggingface.co/K0D3IN/PML-12L-O) • [PML-22L-O](https://huggingface.co/K0D3IN/PML-22L-O) • [PML-6L](https://huggingface.co/K0D3IN/PML-6L)\n\n### Support Open Source AI Research\n**Monero (XMR):** `83iqXtvVu28ZiL9bsATMerSgbFFiD1J1jc96CcxJLEnAW3KBmBKedWnUAeLvLvEA9aBiUBpHQJs1iNHYtkTLZbNUEymobSS`\n\n**Bitcoin (BTC):** `bc1qmnlvpukcgl0hsr7nje0x8555mhtxjt80wtmlxm`") | |
| if __name__ == "__main__": | |
| demo.queue() | |
| demo.launch(server_name="0.0.0.0", theme=gr.themes.Soft()) | |