Model without scatter and with tree

README.md

@@ -99,8 +99,11 @@ def generate_text(model_path, prompt, tokenizer, max_gen_tokens, total_sequence,
     inputs_dict = {}
     inputs_dict['input_ids'] = actual_input[:, :window].reshape(1, window).numpy()
     inputs_dict['attention_mask'] = first_attention
+    index_pos = sum(first_attention[0])
+    inputs_dict['position_ids'] = np.concatenate((np.zeros([1, total_sequence - index_pos], dtype = 'int64'), np.arange(index_pos, dtype = 'int64').reshape(1, index_pos)), axis=1)
+    inputs_dict['tree_attention'] = np.triu(-65504*np.ones(total_sequence), k= 1).astype('float16').reshape(1, 1, total_sequence, total_sequence)
     for name in inputs_names:
-        if name == 'input_ids' or name == 'attention_mask': continue
+        if name == 'input_ids' or name == 'attention_mask' or name == 'position_ids' or name == 'tree_attention': continue
         inputs_dict[name] = np.zeros([1, n_heads, context-window, 128], dtype="float16")
     index = 0
     new_token = np.array([10])
@@ -132,6 +135,9 @@ def generate_text(model_path, prompt, tokenizer, max_gen_tokens, total_sequence,
                 inputs_dict['input_ids']= total_input[:, j:next_index].reshape(1, window)
             elif name == 'attention_mask':
                 inputs_dict['attention_mask'] = np.concatenate((np.zeros((1, total_sequence-next_index), dtype = 'int64'), np.ones((1, next_index), dtype = 'int64')), axis=1)
+            elif name == 'position_ids':
+                inputs_dict['position_ids'] = np.concatenate((np.zeros([1, total_sequence - next_index], dtype = 'int64'), np.arange(next_index, dtype = 'int64').reshape(1, next_index)), axis=1)
+            elif name == 'tree_attention': continue
             else:
                 old_name = name.replace("past_key_values", "present")
                 inputs_dict[name] = outs_dictionary[old_name][:, :, next_index-old_j:context-window+(next_index - old_j), :]

model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3f652af79679d34d3b59b1bfcc4a385bb8bbe37e117d224302512b8d1c576bce
-size 1947171
+oid sha256:9fb4dc57c626c6f6164d2d98e2f15de289093ddc0f74fc006726ef2fb856a0bd
+size 1934549