(model_name, batch_size=8)
| 38 | |
| 39 | |
| 40 | def prepare_text_inputs(model_name, batch_size=8): |
| 41 | print( |
| 42 | f"HF_DATASETS_OFFLINE inside function: {datasets.config.HF_DATASETS_OFFLINE}" |
| 43 | ) |
| 44 | if model_name == "BertForQuestionAnswering" or model_name == "RobertaForQuestionAnswering": |
| 45 | squad_dataset_root = str(llm_datasets_root( |
| 46 | )) + "/" if datasets.config.HF_DATASETS_OFFLINE else "" |
| 47 | squad_dataset_path = squad_dataset_root + "squad_v2" |
| 48 | squad_dataset = load_dataset(squad_dataset_path, trust_remote_code=True) |
| 49 | val_dataset = squad_dataset["validation"] |
| 50 | samples = val_dataset.select(range(batch_size)) |
| 51 | |
| 52 | qa_real_test_inputs = { |
| 53 | 'text': samples["question"], |
| 54 | 'text_pair': samples["context"] |
| 55 | } |
| 56 | return qa_real_test_inputs |
| 57 | elif model_name == "BertForSequenceClassification" or model_name == "RobertaForSequenceClassification": |
| 58 | yelp_dataset_root = str(llm_datasets_root( |
| 59 | )) + "/" if datasets.config.HF_DATASETS_OFFLINE else "fancyzhx/" |
| 60 | yelp_dataset_path = yelp_dataset_root + "yelp_polarity" |
| 61 | yelp_dataset = load_dataset(yelp_dataset_path, trust_remote_code=True) |
| 62 | val_dataset = yelp_dataset["test"] |
| 63 | samples = val_dataset.select(range(batch_size)) |
| 64 | |
| 65 | seqcls_real_test_inputs = {'text': samples['text']} |
| 66 | return seqcls_real_test_inputs |
| 67 | elif model_name == "BertModel" or model_name == "RobertaModel": |
| 68 | #NOTE: For BertModel, it is used as an encoder, so we use dummy input here, |
| 69 | # you can choose whatevert you like, but the numerical accuracy might vary. |
| 70 | test_input = 'To be or not to be: that is the question' |
| 71 | input_strings = [test_input for _ in range(batch_size)] |
| 72 | base_real_test_inputs = {'text': input_strings} |
| 73 | return base_real_test_inputs |
| 74 | |
| 75 | else: |
| 76 | raise NotImplementedError(f"Unknown model {model_name}") |
| 77 | |
| 78 | |
| 79 | def get_engine_name(rank): |
no test coverage detected