(args)
| 359 | |
| 360 | |
| 361 | async def async_main(args): |
| 362 | # Classifier model loading still in main, then passed in |
| 363 | teacher_api_model_name = args.api_model_name_teacher |
| 364 | student_api_model_name = args.api_model_name_student |
| 365 | teacher_api_url = args.teacher_api_url |
| 366 | student_api_url = args.student_api_url |
| 367 | |
| 368 | input_path = args.input_path |
| 369 | output_path = args.output_path |
| 370 | enable_think = args.enable_thinking |
| 371 | classifier_len = args.classifier_len # Use local variable |
| 372 | max_new_tokens_per_sample = args.max_new_tokens |
| 373 | |
| 374 | |
| 375 | system_prompt = args.system_prompt |
| 376 | |
| 377 | dataset = read_jsonl(input_path) |
| 378 | processed_ids = load_processed_ids(output_path) |
| 379 | |
| 380 | llm_tokenizer_teacher = AutoTokenizer.from_pretrained(args.model_name_teacher, trust_remote_code=True) |
| 381 | if llm_tokenizer_teacher.pad_token is None: |
| 382 | llm_tokenizer_teacher.pad_token = llm_tokenizer_teacher.eos_token |
| 383 | |
| 384 | llm_tokenizer_student = AutoTokenizer.from_pretrained(args.model_name_student, trust_remote_code=True) |
| 385 | if llm_tokenizer_student.pad_token is None: |
| 386 | llm_tokenizer_student.pad_token = llm_tokenizer_student.eos_token |
| 387 | |
| 388 | # Load teacher classifier tokenizer and model |
| 389 | teacher_token_classifier_tokenizer = AutoTokenizer.from_pretrained(args.teacher_classifier_path, trust_remote_code=True) |
| 390 | teacher_token_classifier_tokenizer.truncation_side = "left" |
| 391 | teacher_classifier_config = AutoConfig.from_pretrained(args.teacher_classifier_path, num_labels=2) |
| 392 | teacher_classifier_model_base = AutoModelForTokenClassification.from_pretrained( |
| 393 | args.teacher_classifier_path, config=teacher_classifier_config, ignore_mismatched_sizes=True, trust_remote_code=True |
| 394 | ) |
| 395 | |
| 396 | print('====================') |
| 397 | print('Loading teacher classifier from:', args.teacher_classifier_path) |
| 398 | print('====================') |
| 399 | |
| 400 | teacher_classifier_model_base = teacher_classifier_model_base.bfloat16().cuda() |
| 401 | teacher_token_classifier_model = nn.DataParallel(teacher_classifier_model_base).eval() |
| 402 | |
| 403 | # Load student classifier tokenizer and model |
| 404 | student_token_classifier_tokenizer = AutoTokenizer.from_pretrained(args.student_classifier_path, trust_remote_code=True) |
| 405 | student_token_classifier_tokenizer.truncation_side = "left" |
| 406 | student_classifier_config = AutoConfig.from_pretrained(args.student_classifier_path, num_labels=2) |
| 407 | student_classifier_model_base = AutoModelForTokenClassification.from_pretrained( |
| 408 | args.student_classifier_path, config=student_classifier_config, ignore_mismatched_sizes=True, trust_remote_code=True |
| 409 | ) |
| 410 | student_classifier_model_base = student_classifier_model_base.bfloat16().cuda() |
| 411 | print('====================') |
| 412 | print('Loading student classifier from:', args.student_classifier_path) |
| 413 | print('====================') |
| 414 | student_token_classifier_model = nn.DataParallel(student_classifier_model_base).eval() |
| 415 | |
| 416 | |
| 417 | |
| 418 | dataset_queue = deque(dataset) # Use deque to support queue tail append |
no test coverage detected