Get top-n few-shots by mode
(mode, train, question, entities, sql, model, n_shots)
| 61 | |
| 62 | |
| 63 | def get_shots_by_mode(mode, train, question, entities, sql, model, n_shots): |
| 64 | """ |
| 65 | Get top-n few-shots by mode |
| 66 | """ |
| 67 | shots = [] |
| 68 | if mode == 'random': |
| 69 | shots = get_random_samples(train, n_samples=n_shots) |
| 70 | elif mode == 'ques_sim': |
| 71 | shots = get_similar_text_embed_by_knn(train, question, model=model, n_samples=n_shots) |
| 72 | elif mode == 'masked_ques_sim': |
| 73 | mask_question = mask_query_by_entities(question, entities) |
| 74 | shots = get_similar_text_embed_by_knn(train, mask_question, model=model, n_samples=n_shots) |
| 75 | elif mode == 'query_sim': |
| 76 | shots = get_similar_text_embed_by_knn(train, sql, model=model, n_samples=n_shots) |
| 77 | |
| 78 | return shots |
| 79 | |
| 80 | |
| 81 | def get_train_set(question_type: List[str], mode='random', model='transformer', |
no test coverage detected