update model

evaluation.ipynb

@@ -86,14 +86,18 @@
     "    input_ids = pad(input_ids, (0, pad_len), value=1)\n",
     "    ort_inputs = {\n",
     "        'input_ids': input_ids.detach().cpu().numpy(),\n",
-    "        'attention_mask': torch.ones(input_ids.shape).detach().cpu().numpy().astype('int64')\n",
+    "        'attention_mask': torch.cat([torch.ones(input_ids.shape), torch.ones([1, 1])], dim=-1).detach().cpu().numpy().astype('int64')\n",
     "    }\n",
+    "    for i in range(28):\n",
+    "        ort_inputs[\"past_key_values.{}.key\".format(i)] = np.zeros((1,16,1,256), dtype='float32')\n",
+    "        ort_inputs[\"past_key_values.{}.value\".format(i)] = np.zeros((1,16,1,256), dtype='float32')\n",
     "    predictions = session.run(None, ort_inputs)\n",
     "    outputs = torch.from_numpy(predictions[0]) \n",
     "    last_token_logits = outputs[:, -2 - pad_len, :]\n",
     "    pred = last_token_logits.argmax(dim=-1)\n",
     "    total += label.size(0)\n",
     "    hit += (pred == label).sum().item()\n",
+    "\n",
     "acc = hit / total\n",
     "print('acc: ', acc)"
    ]
@@ -132,19 +136,59 @@
     "\n",
     "print(\"prompt: \", prompt)\n",
     "\n",
+    "total_time = 0.0\n",
+    "num_iter = 10\n",
+    "num_warmup = 3\n",
+    "\n",
     "# start\n",
-    "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
-    "for i in range(32):\n",
+    "for idx in range(num_iter):\n",
+    "    text = []\n",
+    "    tic = time.time()\n",
+    "\n",
+    "    input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
+    "\n",
+    "    attention_mask = torch.ones(input_ids.shape[1] +1)\n",
+    "    attention_mask[0] = 0\n",
+    "    attention_mask = attention_mask.unsqueeze(0)\n",
+    "\n",
     "    inp = {'input_ids': input_ids.detach().cpu().numpy(),\n",
-    "            'attention_mask': torch.ones(input_ids.shape).detach().cpu().numpy().astype('int64')}\n",
-    "    output = session.run(None, inp)\n",
-    "    logits = output[0]\n",
-    "    logits = torch.from_numpy(logits)\n",
-    "    next_token_logits = logits[:, -1, :]\n",
-    "    probs = torch.nn.functional.softmax(next_token_logits, dim=-1)\n",
-    "    next_tokens = torch.argmax(probs, dim=-1)\n",
-    "    input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)\n",
-    "print(tokenizer.decode(input_ids[0]))"
+    "            'attention_mask': attention_mask.detach().cpu().numpy().astype('int64')}\n",
+    "    for i in range(28):\n",
+    "        inp[\"past_key_values.{}.key\".format(i)] = torch.zeros([1,16,1,256]).detach().cpu().numpy()\n",
+    "        inp[\"past_key_values.{}.value\".format(i)] = torch.zeros([1,16,1,256]).detach().cpu().numpy()\n",
+    "\n",
+    "    for i in range(32):\n",
+    "\n",
+    "        output = session.run(None, inp)\n",
+    "        logits = output[0]\n",
+    "        logits = torch.from_numpy(logits)\n",
+    "        next_token_logits = logits[:, -1, :]\n",
+    "        probs = torch.nn.functional.softmax(next_token_logits, dim=-1)\n",
+    "        next_tokens = torch.argmax(probs, dim=-1)\n",
+    "        present_kv = output[1]\n",
+    "        for i in range(28):\n",
+    "\n",
+    "            if step == 0:\n",
+    "                inp[\"past_key_values.{}.key\".format(i)] = output[2*i+1][:, :, 1:, :]\n",
+    "                inp[\"past_key_values.{}.value\".format(i)] = output[2*i+2][:, :, 1:, :]\n",
+    "            else:\n",
+    "                inp[\"past_key_values.{}.key\".format(i)] = output[2*i+1]\n",
+    "                inp[\"past_key_values.{}.value\".format(i)] = output[2*i+2]\n",
+    "\n",
+    "        input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)\n",
+    "        if step == 0:\n",
+    "            attention_mask = torch.cat([attention_mask[:, 1:], torch.ones([1, 1])], dim=-1)\n",
+    "        else:\n",
+    "            attention_mask = torch.cat([attention_mask, torch.ones([1, 1])], dim=-1)\n",
+    "\n",
+    "        inp['attention_mask'] = attention_mask.detach().cpu().numpy().astype('int64')\n",
+    "        inp['input_ids'] = input_ids[:, -1:].detach().cpu().numpy()\n",
+    "\n",
+    "    print(tokenizer.decode(input_ids[0]))\n",
+    "    toc = time.time()\n",
+    "    if idx >= num_warmup:\n",
+    "        total_time += (toc - tic)\n",
+    "print(\"Inference latency: %.3f s.\" % (total_time / (num_iter - num_warmup)))"
    ]
   }
  ],

model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:246a7a41aa525a327bf019840ff1f726c7562681b03fc09b9aec5633695fbd95
-size 5852363
+oid sha256:4638154e78188303e3819fff286a65547673a84d5938beec68d3feb8676508c6
+size 6173261

weights.pb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:747ec4484853a81ae90ac3a75518f2dc4623a88b06db69300d961b69166c3087
-size 6790222720
+oid sha256:fcdc4899701755241dfb65b06569c7b5e5d3a6be67c2054f7de431fd6bb1cd48
+size 6057661312