(self, main_model, main_model_tokenizer, assist_model_tokenizer, information_dict,
learning_rate, learning_epochs_nums, result_save_dir,
ensemble_model_output_ids_queue,
assist_model_score_queue_list, main_model_probability_transfer_matrix_list,
assist_model_probability_transfer_matrix_list, device, device_compute, early_stop_string_list=None)
| 10 | |
| 11 | class MainModelThread(threading.Thread): |
| 12 | def __init__(self, main_model, main_model_tokenizer, assist_model_tokenizer, information_dict, |
| 13 | learning_rate, learning_epochs_nums, result_save_dir, |
| 14 | ensemble_model_output_ids_queue, |
| 15 | assist_model_score_queue_list, main_model_probability_transfer_matrix_list, |
| 16 | assist_model_probability_transfer_matrix_list, device, device_compute, early_stop_string_list=None): |
| 17 | self.model = main_model |
| 18 | self.tokenizer = main_model_tokenizer |
| 19 | self.assist_model_tokenizer = assist_model_tokenizer |
| 20 | self.information_dict = information_dict |
| 21 | self.model_streamer = TextStreamer(self.tokenizer) |
| 22 | self.learning_rate = learning_rate |
| 23 | self.learning_epochs_nums = learning_epochs_nums |
| 24 | self.result_save_dir = result_save_dir |
| 25 | self.ensemble_model_output_ids_queue = ensemble_model_output_ids_queue |
| 26 | self.assist_model_score_queue_list = assist_model_score_queue_list |
| 27 | self.main_model_probability_transfer_matrix_list = main_model_probability_transfer_matrix_list |
| 28 | self.assist_model_probability_transfer_matrix_list = assist_model_probability_transfer_matrix_list |
| 29 | self.device = device |
| 30 | self.device_compute = device_compute |
| 31 | self.early_stop_string_list = early_stop_string_list |
| 32 | |
| 33 | super().__init__() |
| 34 | |
| 35 | def run(self) -> None: |
| 36 | main_model_logits_processor_list = LogitsProcessorList() |
nothing calls this directly
no outgoing calls
no test coverage detected