File size: 2,208 Bytes
4dc5bc6 | 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 | #!/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()
|