| 120 | } |
| 121 | |
| 122 | Status QuantizeEmbeddingVariable(const string& input_prefix, |
| 123 | const string& output_prefix, |
| 124 | const std::vector<string>& names, |
| 125 | const std::vector<string>& quant_names, |
| 126 | const std::vector<string>& scale_names, |
| 127 | const TF_DataType data_type, |
| 128 | const bool is_ev) { |
| 129 | BundleReader reader(Env::Default(), input_prefix); |
| 130 | BundleWriter writer(Env::Default(), output_prefix); |
| 131 | const std::set<string> ev_suffix = { |
| 132 | "-freqs", "-freqs_filtered", "-keys", |
| 133 | "-keys_filtered", "-partition_filter_offset", "-partition_offset", |
| 134 | "-versions", "-versions_filtered", "-values"}; |
| 135 | |
| 136 | for (int idx = 0; idx < names.size(); ++idx) { |
| 137 | Status status; |
| 138 | DataType dtype; |
| 139 | TensorShape shape; |
| 140 | string suffix = is_ev ? "-values" : ""; |
| 141 | string value_name = names[idx] + suffix; |
| 142 | status = reader.LookupDtypeAndShape(value_name, &dtype, &shape); |
| 143 | if (!status.ok()) { |
| 144 | errors::InvalidArgument("Invalid variable name:", value_name); |
| 145 | } |
| 146 | Tensor in_tensor(dtype, shape); |
| 147 | status = reader.Lookup(value_name, &in_tensor); |
| 148 | auto in_data = in_tensor.flat<float>(); |
| 149 | |
| 150 | string quant_name = quant_names[idx] + suffix; |
| 151 | if (data_type == TF_DataType::TF_BFLOAT16) { |
| 152 | ConvertToBF16Value(in_tensor, quant_name, writer); |
| 153 | } else if (data_type == TF_DataType::TF_HALF) { |
| 154 | ConvertToHalfValue(in_tensor, quant_name, writer); |
| 155 | } else if (data_type == TF_DataType::TF_INT8) { |
| 156 | ConvertToInt8Value(in_tensor, quant_name, scale_names[idx], writer); |
| 157 | } else { |
| 158 | errors::InvalidArgument("Unsupported data type:", data_type); |
| 159 | } |
| 160 | if (is_ev) { |
| 161 | for (auto it = ev_suffix.cbegin(); it != ev_suffix.cend(); ++it) { |
| 162 | if (*it == "-values") continue; |
| 163 | string tensor_name = names[idx] + *it; |
| 164 | status = reader.LookupDtypeAndShape(tensor_name, &dtype, &shape); |
| 165 | if (status.ok()) { |
| 166 | Tensor tensor(dtype, shape); |
| 167 | status = reader.Lookup(tensor_name, &tensor); |
| 168 | if (status.ok()) { |
| 169 | writer.Add(quant_names[idx] + *it, tensor); |
| 170 | } |
| 171 | } |
| 172 | } |
| 173 | } |
| 174 | } |
| 175 | |
| 176 | if (is_ev) { |
| 177 | WriteRestVariables(reader, writer, names, ev_suffix); |
| 178 | } else { |
| 179 | WriteRestVariables(reader, writer, names); |
nothing calls this directly
no test coverage detected