first pass complete

app.py

@@ -175,5 +175,7 @@ if __name__=='__main__':
                                                                         'key':[(head_id,16,[0,1]) for head_id in range(64)],
                                                                         'val':[(head_id,16,[0,1]) for head_id in range(64)]}})
             logprobs = F.log_softmax(outputs['logits'], dim = -1)
-            preds = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0][1:-1]]
-            st.write([tokenizer.decode([token]) for token in preds])
+            preds_0 = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[0][1:-1]]
+            preds_1 = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1) for probs in logprobs[1][1:-1]]
+            st.write([tokenizer.decode([token]) for token in preds_0])
+            st.write([tokenizer.decode([token]) for token in preds_1])