Coverage for mlprodict/nb_helper.py: 100%

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 

9 

10 

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

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

13 """ 

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

15 

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

17 :param recursive: display subgraph 

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

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

20 :param add_rt_shapes: add information about the shapes 

21 the runtime was able to find out, 

22 the runtime has to be `'python'` 

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

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

25 :param size: graph size 

26 :param html_size: html size 

27 

28 .. versionchanged:: 0.6 

29 Parameter *runtime* was added. 

30 """ 

31 from .onnxrt import OnnxInference 

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) # pragma: no cover 

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 size = args.size 

75 if size == "": 

76 size = None 

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

78 local=args.local, add_rt_shapes=args.add_rt_shapes, 

79 size=size, html_size=args.html_size) 

80 return res 

81 return None 

82 

83 

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

85 """ 

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

87 

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

89 """ 

90 if ip is None: 

91 from IPython import get_ipython 

92 ip = get_ipython() 

93 ip.register_magics(OnnxNotebook)