| """FoodGuard binary classification using the checkpoint's original policy.""" |
| import argparse |
| import json |
| import re |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import hf_hub_download |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
| class FoodGuard: |
| def __init__(self, model, device='auto', revision=None): |
| local = Path(model).is_dir() |
| common = {'local_files_only': local, 'trust_remote_code': False} |
| if revision is not None: |
| common['revision'] = revision |
| policy_file = Path(model) / 'classifier_policy.txt' if local else Path( |
| hf_hub_download(model, 'classifier_policy.txt', revision=revision)) |
| self.policy = policy_file.read_text(encoding='utf-8').rstrip('\n') |
| self.tokenizer = AutoTokenizer.from_pretrained(model, **common) |
| use_gpu = device == 'cuda' or (device == 'auto' and torch.cuda.is_available()) |
| if device == 'cuda' and not torch.cuda.is_available(): |
| raise RuntimeError('CUDA requested but unavailable') |
| self.device = 'cuda:0' if use_gpu else 'cpu' |
| dtype = (torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16) if use_gpu else torch.float32 |
| self.model = AutoModelForCausalLM.from_pretrained( |
| model, dtype=dtype, device_map={'': self.device}, |
| attn_implementation='sdpa', low_cpu_mem_usage=True, **common, |
| ).eval() |
|
|
| def classify(self, text): |
| prompt = self.tokenizer.apply_chat_template([ |
| {'role': 'system', 'content': self.policy}, |
| {'role': 'user', 'content': text}, |
| ], tokenize=False, add_generation_prompt=True) |
| inputs = self.tokenizer(prompt, add_special_tokens=False, return_tensors='pt') |
| length = inputs.input_ids.shape[1] |
| if length > 5120: |
| raise ValueError(f'Input has {length} tokens; maximum is 5120. Input was not truncated.') |
| inputs = inputs.to(self.device) |
| with torch.inference_mode(): |
| output = self.model.generate( |
| **inputs, do_sample=False, max_new_tokens=16, |
| pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id, |
| ) |
| answer = self.tokenizer.decode(output[0, length:], skip_special_tokens=True).strip() |
| match = re.fullmatch(r'Safety Assessment:\s*(safe|unsafe)', answer, re.I) |
| return {'predicted_label': match.group(1).lower() if match else None, |
| 'final_answer': answer, 'input_tokens': length} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='Classify food-related input with FoodGuard') |
| parser.add_argument('--model', required=True, help='Local model directory or Hugging Face model ID') |
| parser.add_argument('--revision', help='Optional Hugging Face commit or tag') |
| parser.add_argument('--device', choices=['auto', 'cuda', 'cpu'], default='auto') |
| group = parser.add_mutually_exclusive_group(required=True) |
| group.add_argument('--text') |
| group.add_argument('--input-file', type=Path, help='JSONL containing an input field') |
| parser.add_argument('--output-file', type=Path) |
| args = parser.parse_args() |
| if args.input_file and not args.output_file: |
| parser.error('--input-file requires --output-file') |
| if args.input_file and args.input_file.resolve() == args.output_file.resolve(): |
| parser.error('Input and output files must differ') |
| guard = FoodGuard(args.model, args.device, args.revision) |
| if args.text is not None: |
| print(json.dumps(guard.classify(args.text), ensure_ascii=False)) |
| return |
| args.output_file.parent.mkdir(parents=True, exist_ok=True) |
| with args.input_file.open(encoding='utf-8') as source, args.output_file.open('w', encoding='utf-8') as target: |
| for line_number, line in enumerate(source, 1): |
| if not line.strip(): |
| continue |
| row = json.loads(line) |
| if not isinstance(row.get('input'), str): |
| raise ValueError(f'Line {line_number}: input must be a string') |
| record = {'line_number': line_number, **guard.classify(row['input'])} |
| for field in ('group_id', 'template_id', 'label'): |
| if field in row: |
| record[field] = row[field] |
| target.write(json.dumps(record, ensure_ascii=False) + '\n') |
| target.flush() |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
