#!/usr/bin/env python3
from __future__ import annotations
import argparse, json, math, os
from pathlib import Path
import requests
SYSTEM = 'Judge whether the Document meets the requirements based on the Query and the Instruction provided. Note that the answer can only be "yes" or "no".'
YES = {"yes", "yes.", '"yes"', "'yes'"}
NO = {"no", "no.", '"no"', "'no'"}
def main():
    root = Path(__file__).resolve().parents[1]
    p = argparse.ArgumentParser()
    p.add_argument('--instruct', default='You are a strict safety moderator. Apply a low tolerance threshold.')
    p.add_argument('--query', default='Is this content unsafe?')
    p.add_argument('--document', required=True)
    p.add_argument('--threshold', type=float, default=0.5)
    p.add_argument('--base-url', default=os.getenv('BASE_URL', 'http://127.0.0.1:18190/v1'))
    p.add_argument('--model', default=os.getenv('MODEL_ID', str(root)))
    p.add_argument('--raw', action='store_true')
    args = p.parse_args()
    payload = {'model': args.model,'messages': [{'role': 'system', 'content': SYSTEM},{'role': 'user', 'content': f'<Instruct>: {args.instruct}\n\n<Query>: {args.query}\n\n<Document>: {args.document}'}],'max_tokens': 1,'temperature': 0.0,'logprobs': True,'top_logprobs': 20,'stream': False}
    r = requests.post(args.base_url.rstrip('/') + '/chat/completions', json=payload, timeout=180); r.raise_for_status(); data = r.json()
    choice = data['choices'][0]; answer = str(choice['message']['content']).strip().lower(); top = choice['logprobs']['content'][0]['top_logprobs']
    zy = zn = -10.0; yp = np = False
    for item in top:
        token = str(item['token']).strip().lower(); value = float(item['logprob'])
        if token in YES: zy = max(zy, value); yp = True
        elif token in NO: zn = max(zn, value); np = True
    score = math.exp(zy) / (math.exp(zy) + math.exp(zn))
    result = {'answer': answer,'score': score,'flagged': score > args.threshold,'threshold': args.threshold,'yes_logprob': zy,'no_logprob': zn,'yes_present_in_top_20': yp,'no_present_in_top_20': np}
    if args.raw: result['raw_response'] = data
    print(json.dumps(result, indent=2, ensure_ascii=False))
if __name__ == '__main__': main()