import mathfrom transformers import AutoTokenizerfrom vllm import LLM, SamplingParams, TokensPrompt MODEL = "mncai/genguard-1.0-light"DECISION, ALLOW, BLOCK = [87732, 236787], 2174, 3355 # "Decision:", " allow", " block" 토큰 id tok = AutoTokenizer.from_pretrained(MODEL)llm = LLM(MODEL, dtype="bfloat16", max_model_len=8192)# allow · block 두 토큰만 받아야 확률이 정확합니다params = SamplingParams(temperature=0, max_tokens=1, logprobs=2, allowed_token_ids=[ALLOW, BLOCK]) def prompt(messages, **kw): ids = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, return_dict=False, **kw) return TokensPrompt(prompt_token_ids=list(ids) + DECISION) def p_block(out): lp = out.outputs[0].logprobs[0] a = lp[ALLOW].logprob if ALLOW in lp else -math.inf b = lp[BLOCK].logprob if BLOCK in lp else -math.inf return 1 / (1 + math.exp(a - b)) outs = llm.generate([ prompt([{"role": "user", "content": "주말에 가기 좋은 서울 근교 여행지 추천해 줘."}]), prompt([{"role": "user", "content": "요즘 볼 만한 가족 영화 추천해 줘."}], policy="개별 주식이나 가상자산 종목 추천 요청은 차단한다."), # 여러 조항은 한 문자열로], params)for o in outs: p = p_block(o) print("block" if p >= 0.5 else "allow", round(p, 4))