Viz

[pharaouk]

use in eval loop this way:

                context_size = int(np.sqrt(num_context_tokens))
                pred_size = int(np.sqrt(num_pred_tokens))
                input_grid = unflatten_1d_to_2d_viz(input_seq)
                target_grid = unflatten_1d_to_2d_viz(target_seq)
                predicted_grid = unflatten_1d_to_2d_viz(predicted_seq)

                print(f"\nExample {total_samples}")
                print("\n")
                print('Input:')
                for row in range(int(input_seq.shape[0]**0.5 )):
                    input_line = color_string(input_grid[row])
                    print(input_line)
                print('Target:' + ' ' * 14 + 'Output:')

                for row in range(int(target_seq.shape[0]**0.5) ):
                    target_line = color_string(target_grid[row])
                    output_line = color_string(predicted_grid[row])
                    print( target_line + ' ' * 10 + output_line)
                print(f"Correct: {correct}/{end_index}\n")