1""" 

2@file 

3@brief Helpers for notebooks. 

4""" 

5from IPython.core.magic import magics_class, line_magic 

6from jyquickhelper import RenderJsDot 

7from pyquickhelper.ipythonhelper import MagicCommandParser, MagicClassWithHelpers 

8from pyquickhelper.cli.cli_helper import create_cli_parser 

9from ..onnx_inference import OnnxInference 

10 

11 

12def onnxview(graph, recursive=False, local=False, add_rt_shapes=False, 

13 runtime='python', size=None, html_size=None): 

14 """ 

15 Displays an :epkg:`ONNX` graph into a notebook. 

16 

17 :param graph: filename, bytes, or :epkg:`onnx` graph. 

18 :param recursive: display subgraph 

19 :param local: use local path to javascript dependencies, 

20 recommanded option if used on :epkg:`MyBinder`) 

21 :param add_rt_shapes: add information about the shapes 

22 the runtime was able to find out, 

23 the runtime has to be `'python'` 

24 :param runtime: the view fails if a runtime does not implement a specific 

25 node unless *runtime* is `'empty'` 

26 :param size: graph size 

27 :param html_size: html size 

28 

29 .. versionchanged:: 0.6 

30 Parameter *runtime* was added. 

31 """ 

32 sess = OnnxInference(graph, skip_run=not add_rt_shapes, runtime=runtime) 

33 dot = sess.to_dot(recursive=recursive, 

34 add_rt_shapes=add_rt_shapes, size=size) 

35 if html_size is not None: 

36 return RenderJsDot(dot, local=local, width=html_size, height=html_size) 

37 return RenderJsDot(dot, local=local) 

38 

39 

40@magics_class 

41class OnnxNotebook(MagicClassWithHelpers): 

42 

43 """ 

44 Defines magic commands to help with notebooks 

45 

46 .. versionadded:: 1.1 

47 """ 

48 

49 @line_magic 

50 def onnxview(self, line): 

51 """ 

52 Defines ``%onnxview`` 

53 which displays an :epkg:`ONNX` graph. 

54 

55 .. nbref:: 

56 :title: onnxview 

57 

58 The magic command ``%onnxview model_onnx`` is equivalent to function 

59 :func:`onnxview <mlprodict.onnxrt.doc.nb_helper.onnxview>`: 

60 

61 :: 

62 

63 onnx_view(model_onnx) 

64 

65 It displays a visual representation of an :epkg:`ONNX` graph. 

66 

67 """ 

68 parser = self.get_parser( 

69 lambda: create_cli_parser(onnxview, cls=MagicCommandParser, 

70 positional={'graph'}), "onnxview") 

71 args = self.get_args(line, parser) 

72 

73 if args is not None: 

74 res = onnxview(args.graph, recursive=args.recursive, 

75 local=args.local, add_rt_shapes=args.add_rt_shapes, 

76 size=args.size, html_size=args.html_size) 

77 return res 

78 return None 

79 

80 

81def register_onnx_magics(ip=None): # pragma: no cover 

82 """ 

83 Register magics function, can be called from a notebook. 

84 

85 @param ip from ``get_ipython()`` 

86 """ 

87 if ip is None: 

88 from IPython import get_ipython 

89 ip = get_ipython() 

90 ip.register_magics(OnnxNotebook)