Obtain train set according to different input parameters
(question_type: List[str], mode='random', model='transformer',
index_version='train_fold1_v1', ques_type_mode='each')
| 79 | |
| 80 | |
| 81 | def get_train_set(question_type: List[str], mode='random', model='transformer', |
| 82 | index_version='train_fold1_v1', ques_type_mode='each'): |
| 83 | """ |
| 84 | Obtain train set according to different input parameters |
| 85 | """ |
| 86 | |
| 87 | output_dir = f'fewshot/index_{index_version}/' |
| 88 | output_file_map = { |
| 89 | 'random': f'query_embed_{model}.pkl', |
| 90 | 'ques_sim': f'query_embed_{model}.pkl', |
| 91 | 'masked_ques_sim': f'masked_query_embed_{model}.pkl', |
| 92 | 'query_sim': f'pred_embed_{model}.pkl', |
| 93 | } |
| 94 | output_file = output_file_map.get(mode, '') |
| 95 | |
| 96 | if not output_file: |
| 97 | return {} |
| 98 | |
| 99 | with open(output_dir + output_file, 'rb') as f: |
| 100 | data = pickle.load(f) |
| 101 | |
| 102 | if ques_type_mode == 'each': |
| 103 | all_train = {q_type: {} for q_type in question_type} |
| 104 | for emb, sample in data.items(): |
| 105 | q_type = sample.get('question_type', '') |
| 106 | if q_type in question_type: |
| 107 | all_train[q_type][emb] = sample |
| 108 | else: |
| 109 | all_train = {} |
| 110 | for emb, sample in data.items(): |
| 111 | q_type = sample.get('question_type', '') |
| 112 | if q_type in question_type: |
| 113 | all_train[emb] = sample |
| 114 | return all_train |
| 115 | |
| 116 | |
| 117 | def get_fewshots(question: str, entities: List[str], sql: str, question_type: List[str], |