MCPcopy Create free account
hub / github.com/OpenNMT/CTranslate2 / MockModel

Method MockModel

tests/attention_test.cc:6–28  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

4class MockModel : public models::Model {
5public:
6 MockModel(dim_t num_heads, dim_t num_heads_kv) {
7 const dim_t d_model = 64;
8 const dim_t d_head = d_model / num_heads;
9
10 std::vector<float> linear_0_data(num_heads * d_head * d_model, 0.01f);
11 std::vector<float> linear_1_data(2 * num_heads_kv * d_head * d_model, 0.01f);
12
13 register_variable("attn/linear_0/weight",
14 StorageView({num_heads * d_head, d_model}, linear_0_data));
15 register_variable("attn/linear_1/weight",
16 StorageView({2 * num_heads_kv * d_head, d_model}, linear_1_data));
17 register_variable("attn/linear_2/weight",
18 StorageView({d_model, num_heads * d_head}, DataType::FLOAT32));
19 register_variable("attn/q_norm/gamma",
20 StorageView({d_model}, std::vector<float>(d_model, 1.0f)));
21 register_variable("attn/k_norm/gamma",
22 StorageView({d_head}, std::vector<float>(d_head, 1.0f)));
23
24 register_variable("attn/num_heads_kv",
25 StorageView(static_cast<int32_t>(num_heads_kv)));
26
27 set_compute_type(ComputeType::FLOAT32, Device::CPU, 0, false);
28 }
29protected:
30 std::unique_ptr<Model> clone() const override { return nullptr; }
31};

Callers

nothing calls this directly

Calls 1

StorageViewClass · 0.85

Tested by

no test coverage detected