Refactor model type / training tables, simplify reverse mapping

PiperOrigin-RevId: 647069372
This commit is contained in:
Paul Chang 2024-06-26 13:58:40 -07:00 committed by Copybara-Service
parent aa57fc3952
commit 78e96fdc70
1 changed files with 33 additions and 26 deletions

View File

@ -28,36 +28,47 @@
namespace gcpp {
constexpr const char* kModelFlags[] = {
"2b-pt", "2b-it", // Gemma 2B
"7b-pt", "7b-it", // Gemma 7B
"gr2b-pt", "gr2b-it", // RecurrentGemma
"tiny", // Gemma Tiny (mostly for debugging)
};
constexpr Model kModelTypes[] = {
Model::GEMMA_2B, Model::GEMMA_2B, // Gemma 2B
Model::GEMMA_7B, Model::GEMMA_7B, // Gemma 7B
Model::GRIFFIN_2B, Model::GRIFFIN_2B, // RecurrentGemma
Model::GEMMA_TINY, // Gemma Tiny
};
constexpr ModelTraining kModelTraining[] = {
ModelTraining::GEMMA_PT, ModelTraining::GEMMA_IT, // Gemma 2B
ModelTraining::GEMMA_PT, ModelTraining::GEMMA_IT, // Gemma 7B
ModelTraining::GEMMA_PT, ModelTraining::GEMMA_IT, // RecurrentGemma
ModelTraining::GEMMA_IT, // Gemma Tiny
};
constexpr size_t kNumModelFlags = std::end(kModelFlags) - std::begin(kModelFlags);
static_assert(kNumModelFlags ==
std::end(kModelTypes) - std::begin(kModelTypes));
static_assert(kNumModelFlags ==
std::end(kModelTraining) - std::begin(kModelTraining));
const char* ParseModelTypeAndTraining(const std::string& model_flag,
Model& model, ModelTraining& training) {
constexpr const char* kModelFlags[] = {
"2b-pt", "7b-pt", "gr2b-pt", "2b-it", "7b-it", "gr2b-it", "tiny",
};
constexpr Model kModelTypes[] = {
Model::GEMMA_2B, Model::GEMMA_7B, Model::GRIFFIN_2B, Model::GEMMA_2B,
Model::GEMMA_7B, Model::GRIFFIN_2B, Model::GEMMA_TINY,
};
constexpr ModelTraining kModelTraining[] = {
ModelTraining::GEMMA_PT, ModelTraining::GEMMA_PT, ModelTraining::GEMMA_PT,
ModelTraining::GEMMA_IT, ModelTraining::GEMMA_IT, ModelTraining::GEMMA_IT,
ModelTraining::GEMMA_IT,
};
constexpr size_t kNum = std::end(kModelFlags) - std::begin(kModelFlags);
static char kErrorMessageBuffer[kNum * 8 + 1024] =
static char kErrorMessageBuffer[kNumModelFlags * 8 + 1024] =
"Invalid or missing model flag, need to specify one of ";
for (size_t i = 0; i + 1 < kNum; i++) {
for (size_t i = 0; i + 1 < kNumModelFlags; i++) {
strcat(kErrorMessageBuffer, kModelFlags[i]); // NOLINT
strcat(kErrorMessageBuffer, ", "); // NOLINT
}
strcat(kErrorMessageBuffer, kModelFlags[kNum - 1]); // NOLINT
strcat(kErrorMessageBuffer, kModelFlags[kNumModelFlags - 1]); // NOLINT
strcat(kErrorMessageBuffer, "."); // NOLINT
std::string model_type_lc = model_flag;
std::transform(begin(model_type_lc), end(model_type_lc), begin(model_type_lc),
[](unsigned char c) { return std::tolower(c); });
for (size_t i = 0; i < kNum; i++) {
for (size_t i = 0; i < kNumModelFlags; i++) {
if (kModelFlags[i] == model_type_lc) {
model = kModelTypes[i];
training = kModelTraining[i];
@ -69,14 +80,10 @@ const char* ParseModelTypeAndTraining(const std::string& model_flag,
}
const char* ModelString(Model model, ModelTraining training) {
if (model == Model::GEMMA_TINY) return "tiny";
static_assert(static_cast<size_t>(ModelTraining::GEMMA_IT) == 0);
constexpr const char* k2B[] = {"2b-it", "2b-pt"};
constexpr const char* k7B[] = {"7b-it", "7b-pt"};
constexpr const char* kGr2B[] = {"gr2b-it", "gr2b-pt"};
if (model == Model::GEMMA_2B) return k2B[static_cast<size_t>(training)];
if (model == Model::GEMMA_7B) return k7B[static_cast<size_t>(training)];
if (model == Model::GRIFFIN_2B) return kGr2B[static_cast<size_t>(training)];
for (size_t i = 0; i < kNumModelFlags; i++) {
if (kModelTypes[i] == model && kModelTraining[i] == training)
return kModelFlags[i];
}
HWY_ABORT("Unknown model %d training %d\n", static_cast<int>(model),
static_cast<int>(training));
}