Get fewshots according to different input parameters :param question: the question :param entities: the entities :param sql: the sql :param question_type: the question type :param n_shots: the num of shots :param mode: the retrieval mode, which can be random, ques_sim,
(question: str, entities: List[str], sql: str, question_type: List[str],
n_shots=1, mode='random', model='transformer',
index_version='train_fold1_v1', ques_type_mode='each')
| 115 | |
| 116 | |
| 117 | def get_fewshots(question: str, entities: List[str], sql: str, question_type: List[str], |
| 118 | n_shots=1, mode='random', model='transformer', |
| 119 | index_version='train_fold1_v1', ques_type_mode='each') -> Union[Dict, List]: |
| 120 | """ |
| 121 | Get fewshots according to different input parameters |
| 122 | |
| 123 | :param question: the question |
| 124 | :param entities: the entities |
| 125 | :param sql: the sql |
| 126 | :param question_type: the question type |
| 127 | :param n_shots: the num of shots |
| 128 | :param mode: the retrieval mode, which can be random, ques_sim, masked_ques_sim, query_sim |
| 129 | :param model: the embedding model, which can be transformer or openai |
| 130 | :param index_version: the version name of index library |
| 131 | :param ques_type_mode: the question type recall mode, each for separate recall, all is combined recall |
| 132 | :return: the fewshot samples |
| 133 | -ques_type_mode=each |
| 134 | - Type: Dict |
| 135 | - key is the question type |
| 136 | - value is a list of few shots recalled corresponding to the question type, each element is a sample, and each sample is a Dict |
| 137 | -ques_type_mode=all |
| 138 | - Type: List |
| 139 | - Each element is a sample and each sample is a Dict |
| 140 | """ |
| 141 | |
| 142 | assert mode in ['random', 'ques_sim', 'masked_ques_sim', 'query_sim'] |
| 143 | assert model in ['transformer', 'openai'] |
| 144 | assert ques_type_mode in ['each', 'all'] |
| 145 | |
| 146 | all_train = get_train_set(question_type, mode, model, index_version, ques_type_mode) |
| 147 | |
| 148 | if ques_type_mode == 'each': |
| 149 | all_shots = {} |
| 150 | for q_type in question_type: |
| 151 | train = all_train[q_type] |
| 152 | shots = get_shots_by_mode(mode, train, question, entities, sql, model, n_shots) |
| 153 | all_shots[q_type] = shots |
| 154 | else: |
| 155 | all_shots = get_shots_by_mode(mode, all_train, question, entities, sql, model, n_shots) |
| 156 | |
| 157 | return all_shots |
| 158 | |
| 159 | |
| 160 | if __name__ == '__main__': |
no test coverage detected