(args, input_str, tokenizer, lm_model, device, logger)
| 841 | |
| 842 | |
| 843 | def get_model_output(args, input_str, tokenizer, lm_model, device, logger): |
| 844 | input_ids = tokenizer(input_str, return_tensors="pt", max_length=args["max_input_len"] , truncation=True).input_ids |
| 845 | |
| 846 | sample_outputs = lm_model.generate( |
| 847 | input_ids.to(device), |
| 848 | max_length=16, |
| 849 | num_return_sequences=args['beams'], |
| 850 | num_beams=args['beams'], |
| 851 | ) |
| 852 | |
| 853 | lm_pred = sample_outputs |
| 854 | |
| 855 | # Take the first prediction that is not "look around" |
| 856 | logger.info("Top N Predictions:") |
| 857 | predStrs = [] |
| 858 | for i, pred in enumerate(lm_pred): |
| 859 | text = tokenizer.decode(pred) |
| 860 | text = post_process_generation(text) |
| 861 | logger.info("\t" + str(i) + "\t" + str(text) ) |
| 862 | predStrs.append(text) |
| 863 | |
| 864 | return predStrs |
| 865 | |
| 866 | |
| 867 | def post_process_generation(raw_pred): |
nothing calls this directly
no test coverage detected