File size: 4,721 Bytes
44810a8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """
Example: End-to-End Generation with GCLM LogitsProcessor
Demonstrates strict budget JSON completion and dead-end avoidance with Hugging Face Transformers.
"""
import argparse
import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList
from core.fsm_builder import ReachabilityFSM
from core.logit_processor import GoalReachabilityLogitsProcessor
from core.compiler import FSMCompiler
def run_demo(model_name: str = "Qwen/Qwen2.5-0.5B", budget: int = 15):
print(f"\n[DEMO] Loading Tokenizer and Model: {model_name}...")
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
device_map="auto" if device == "cuda" else None,
)
if device == "cpu":
model.to("cpu")
except Exception as e:
print(f"[WARN] Could not load remote model ({e}). Using GPT-2 fallback or mock...")
try:
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
except Exception:
print("[WARN] Hugging Face offline or unavailable. Running in synthetic mode.")
return
vocab_size = model.config.vocab_size
compiler = FSMCompiler(vocab_size=vocab_size, tokenizer=tokenizer, device=device)
# 1. Define JSON structure with optional fields:
# {"status": "ok", "code": 200}
# Open: '{"' or '{'
# Close: '}'
# We want model to complete valid JSON before budget runs out.
prompt = "Generate a JSON response for server status: "
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
# Compile a simple strict JSON FSM
open_tokens = tokenizer.encode('{"', add_special_tokens=False)
kv1_tokens = tokenizer.encode('status":"ok",', add_special_tokens=False)
kv2_tokens = tokenizer.encode('status":"ok"', add_special_tokens=False)
kv3_tokens = tokenizer.encode('code":200', add_special_tokens=False)
close_tokens = tokenizer.encode('}', add_special_tokens=False)
eos_id = tokenizer.eos_token_id or 0
# Build FSM
# 0: Start -> 1: Open
# 1 -> 2: status:"ok", -> 1 (loop)
# 1 -> 3: status:"ok" -> 4: Close
# 1 -> 5: code":200 -> 4: Close
# 1 -> 4: Close
# 4 -> 6: Goal (on EOS)
fsm = ReachabilityFSM(num_states=10, vocab_size=vocab_size, device=device)
# 0 -> 1 on open
for t in open_tokens:
fsm.add_transition(0, t, 1)
# 1 -> 4 on close
for t in close_tokens:
fsm.add_transition(1, t, 4)
# 1 -> loop or 1 -> close via KV
if len(kv1_tokens) > 0:
fsm.add_transition(1, kv1_tokens[0], 1)
if len(kv2_tokens) > 0:
fsm.add_transition(1, kv2_tokens[0], 4)
if len(kv3_tokens) > 0:
fsm.add_transition(1, kv3_tokens[0], 4)
# 4 -> 6 on eos
fsm.add_transition(4, eos_id, 6)
fsm.add_transition(6, eos_id, 6) # goal self-loop
fsm.set_goal_states([6])
fsm.build_reachability(max_steps=budget)
# Setup GCLM LogitsProcessor
gclm_processor = GoalReachabilityLogitsProcessor(fsm=fsm, max_budget=budget)
logits_processors = LogitsProcessorList([gclm_processor])
print(f"\nPrompt: '{prompt}'")
print(f"Token Budget: {budget} tokens\n")
# Generate with GCLM
print("--- 1. Generation with GCLM (Guaranteed Completion) ---")
output_gclm = model.generate(
input_ids,
max_new_tokens=budget,
logits_processor=logits_processors,
do_sample=True,
temperature=0.7,
pad_token_id=eos_id,
)
generated_text_gclm = tokenizer.decode(output_gclm[0], skip_special_tokens=False)
print(f"Output:\n{generated_text_gclm}\n")
# Generate with Vanilla (Unconstrained)
print("--- 2. Generation with Vanilla (Unconstrained) ---")
output_vanilla = model.generate(
input_ids,
max_new_tokens=budget,
do_sample=True,
temperature=0.7,
pad_token_id=eos_id,
)
generated_text_vanilla = tokenizer.decode(output_vanilla[0], skip_special_tokens=False)
print(f"Output:\n{generated_text_vanilla}\n")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run GCLM generation demo.")
parser.add_argument("--model", type=str, default="Qwen/Qwen2.5-0.5B", help="Model name or path")
parser.add_argument("--budget", type=int, default=15, help="Max new tokens budget")
args = parser.parse_args()
run_demo(model_name=args.model, budget=args.budget)
|