(question, episode)
| 409 | return res |
| 410 | |
| 411 | def run_episode(question, episode) -> Dict: |
| 412 | |
| 413 | global _llm, _tokenizer |
| 414 | |
| 415 | raw_text, _ = make_context( |
| 416 | tokenizer = _tokenizer, |
| 417 | query = question, |
| 418 | system = "You are a helpful assistant.", |
| 419 | max_window_size = 6144, |
| 420 | chat_format = "chatml") |
| 421 | |
| 422 | input_ids = _tokenizer(raw_text, return_tensors='pt', padding='longest') |
| 423 | |
| 424 | attention_mask = input_ids.attention_mask |
| 425 | input_ids = input_ids.input_ids |
| 426 | |
| 427 | input_ids = input_ids.to(_llm.device) |
| 428 | attention_mask = attention_mask.to(_llm.device) |
| 429 | |
| 430 | out_ids = _llm.generate( |
| 431 | input_ids = input_ids, |
| 432 | attention_mask = attention_mask, |
| 433 | do_sample = False, |
| 434 | num_beams = 1, |
| 435 | length_penalty = 1, |
| 436 | num_return_sequences = 1, |
| 437 | use_cache = True, |
| 438 | pad_token_id = _tokenizer.eod_id, |
| 439 | eos_token_id = _tokenizer.eod_id, |
| 440 | min_new_tokens = 1, |
| 441 | max_new_tokens = 30, |
| 442 | ) |
| 443 | |
| 444 | padding_len = input_ids[0].eq(_tokenizer.pad_token_id).sum().item() |
| 445 | response = decode_tokens( |
| 446 | out_ids[0][padding_len: ], |
| 447 | _tokenizer, |
| 448 | raw_text_len = len(raw_text), |
| 449 | context_length = input_ids.size(1) - padding_len, |
| 450 | chat_format = "chatml", |
| 451 | verbose = False, |
| 452 | errors = 'replace' |
| 453 | ) |
| 454 | |
| 455 | episode["pred"] = mapping_actions(response) |
| 456 | |
| 457 | return episode |
| 458 | |
| 459 | |
| 460 | def predict(args): |
nothing calls this directly
no test coverage detected