| |
| 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() |
|
|
