| |
| |
| <!DOCTYPE html> |
| <!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]--> |
| <!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]--> |
| <head> |
| <meta charset="utf-8"> |
| |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> |
| |
| <title>apache_beam.ml.transforms.base — Apache Beam 2.55.1 documentation</title> |
| |
| |
| |
| |
| |
| |
| |
| |
| <script type="text/javascript" src="../../../../_static/js/modernizr.min.js"></script> |
| |
| |
| <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script> |
| <script type="text/javascript" src="../../../../_static/jquery.js"></script> |
| <script type="text/javascript" src="../../../../_static/underscore.js"></script> |
| <script type="text/javascript" src="../../../../_static/doctools.js"></script> |
| <script type="text/javascript" src="../../../../_static/language_data.js"></script> |
| <script async="async" type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script> |
| |
| <script type="text/javascript" src="../../../../_static/js/theme.js"></script> |
| |
| |
| |
| |
| <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" /> |
| <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" /> |
| <link rel="index" title="Index" href="../../../../genindex.html" /> |
| <link rel="search" title="Search" href="../../../../search.html" /> |
| </head> |
| |
| <body class="wy-body-for-nav"> |
| |
| |
| <div class="wy-grid-for-nav"> |
| |
| <nav data-toggle="wy-nav-shift" class="wy-nav-side"> |
| <div class="wy-side-scroll"> |
| <div class="wy-side-nav-search" > |
| |
| |
| |
| <a href="../../../../index.html" class="icon icon-home"> Apache Beam |
| |
| |
| |
| </a> |
| |
| |
| |
| |
| <div class="version"> |
| 2.55.1 |
| </div> |
| |
| |
| |
| |
| <div role="search"> |
| <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get"> |
| <input type="text" name="q" placeholder="Search docs" /> |
| <input type="hidden" name="check_keywords" value="yes" /> |
| <input type="hidden" name="area" value="default" /> |
| </form> |
| </div> |
| |
| |
| </div> |
| |
| <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation"> |
| |
| |
| |
| |
| |
| |
| <ul> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.coders.html">apache_beam.coders package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.dataframe.html">apache_beam.dataframe package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.io.html">apache_beam.io package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.metrics.html">apache_beam.metrics package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.ml.html">apache_beam.ml package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.options.html">apache_beam.options package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.portability.html">apache_beam.portability package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.runners.html">apache_beam.runners package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.testing.html">apache_beam.testing package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.transforms.html">apache_beam.transforms package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.typehints.html">apache_beam.typehints package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.utils.html">apache_beam.utils package</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.yaml.html">apache_beam.yaml package</a></li> |
| </ul> |
| <ul> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.error.html">apache_beam.error module</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.pipeline.html">apache_beam.pipeline module</a></li> |
| <li class="toctree-l1"><a class="reference internal" href="../../../../apache_beam.pvalue.html">apache_beam.pvalue module</a></li> |
| </ul> |
| |
| |
| |
| </div> |
| </div> |
| </nav> |
| |
| <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"> |
| |
| |
| <nav class="wy-nav-top" aria-label="top navigation"> |
| |
| <i data-toggle="wy-nav-top" class="fa fa-bars"></i> |
| <a href="../../../../index.html">Apache Beam</a> |
| |
| </nav> |
| |
| |
| <div class="wy-nav-content"> |
| |
| <div class="rst-content"> |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| <div role="navigation" aria-label="breadcrumbs navigation"> |
| |
| <ul class="wy-breadcrumbs"> |
| |
| <li><a href="../../../../index.html">Docs</a> »</li> |
| |
| <li><a href="../../../index.html">Module code</a> »</li> |
| |
| <li>apache_beam.ml.transforms.base</li> |
| |
| |
| <li class="wy-breadcrumbs-aside"> |
| |
| </li> |
| |
| </ul> |
| |
| |
| <hr/> |
| </div> |
| <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article"> |
| <div itemprop="articleBody"> |
| |
| <h1>Source code for apache_beam.ml.transforms.base</h1><div class="highlight"><pre> |
| <span></span><span class="c1">#</span> |
| <span class="c1"># Licensed to the Apache Software Foundation (ASF) under one or more</span> |
| <span class="c1"># contributor license agreements. See the NOTICE file distributed with</span> |
| <span class="c1"># this work for additional information regarding copyright ownership.</span> |
| <span class="c1"># The ASF licenses this file to You under the Apache License, Version 2.0</span> |
| <span class="c1"># (the "License"); you may not use this file except in compliance with</span> |
| <span class="c1"># the License. You may obtain a copy of the License at</span> |
| <span class="c1">#</span> |
| <span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span> |
| <span class="c1">#</span> |
| <span class="c1"># Unless required by applicable law or agreed to in writing, software</span> |
| <span class="c1"># distributed under the License is distributed on an "AS IS" BASIS,</span> |
| <span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span> |
| <span class="c1"># See the License for the specific language governing permissions and</span> |
| <span class="c1"># limitations under the License.</span> |
| |
| <span class="kn">import</span> <span class="nn">abc</span> |
| <span class="kn">import</span> <span class="nn">collections</span> |
| <span class="kn">import</span> <span class="nn">logging</span> |
| <span class="kn">import</span> <span class="nn">os</span> |
| <span class="kn">import</span> <span class="nn">tempfile</span> |
| <span class="kn">import</span> <span class="nn">uuid</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Generic</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Mapping</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Optional</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Sequence</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">TypeVar</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span> |
| |
| <span class="kn">import</span> <span class="nn">jsonpickle</span> |
| <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span> |
| |
| <span class="kn">import</span> <span class="nn">apache_beam</span> <span class="k">as</span> <span class="nn">beam</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.io.filesystems</span> <span class="kn">import</span> <span class="n">FileSystems</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.metrics.metric</span> <span class="kn">import</span> <span class="n">Metrics</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.ml.inference.base</span> <span class="kn">import</span> <span class="n">ModelHandler</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.ml.inference.base</span> <span class="kn">import</span> <span class="n">ModelT</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.ml.inference.base</span> <span class="kn">import</span> <span class="n">RunInferenceDLQ</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.options.pipeline_options</span> <span class="kn">import</span> <span class="n">PipelineOptions</span> |
| |
| <span class="n">_LOGGER</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span> |
| <span class="n">_ATTRIBUTE_FILE_NAME</span> <span class="o">=</span> <span class="s1">'attributes.json'</span> |
| |
| <span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span> |
| <span class="s1">'MLTransform'</span><span class="p">,</span> |
| <span class="s1">'ProcessHandler'</span><span class="p">,</span> |
| <span class="s1">'MLTransformProvider'</span><span class="p">,</span> |
| <span class="s1">'BaseOperation'</span><span class="p">,</span> |
| <span class="s1">'EmbeddingsManager'</span> |
| <span class="p">]</span> |
| |
| <span class="n">TransformedDatasetT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'TransformedDatasetT'</span><span class="p">)</span> |
| <span class="n">TransformedMetadataT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'TransformedMetadataT'</span><span class="p">)</span> |
| |
| <span class="c1"># Input/Output types to the MLTransform.</span> |
| <span class="n">MLTransformOutputT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'MLTransformOutputT'</span><span class="p">)</span> |
| <span class="n">ExampleT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'ExampleT'</span><span class="p">)</span> |
| |
| <span class="c1"># Input to the apply() method of BaseOperation.</span> |
| <span class="n">OperationInputT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'OperationInputT'</span><span class="p">)</span> |
| <span class="c1"># Output of the apply() method of BaseOperation.</span> |
| <span class="n">OperationOutputT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">'OperationOutputT'</span><span class="p">)</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_convert_list_of_dicts_to_dict_of_lists</span><span class="p">(</span> |
| <span class="n">list_of_dicts</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]])</span> <span class="o">-></span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">]]:</span> |
| <span class="n">keys_to_element_list</span> <span class="o">=</span> <span class="n">collections</span><span class="o">.</span><span class="n">defaultdict</span><span class="p">(</span><span class="nb">list</span><span class="p">)</span> |
| <span class="n">input_keys</span> <span class="o">=</span> <span class="n">list_of_dicts</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> |
| <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="n">list_of_dicts</span><span class="p">:</span> |
| <span class="k">if</span> <span class="nb">set</span><span class="p">(</span><span class="n">d</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> <span class="o">!=</span> <span class="nb">set</span><span class="p">(</span><span class="n">input_keys</span><span class="p">):</span> |
| <span class="n">extra_keys</span> <span class="o">=</span> <span class="nb">set</span><span class="p">(</span><span class="n">d</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> <span class="o">-</span> <span class="nb">set</span><span class="p">(</span><span class="n">input_keys</span><span class="p">)</span> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span> |
| <span class="n">d</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> <span class="o">></span> <span class="nb">len</span><span class="p">(</span><span class="n">input_keys</span><span class="p">)</span> <span class="k">else</span> <span class="nb">set</span><span class="p">(</span><span class="n">input_keys</span><span class="p">)</span> <span class="o">-</span> <span class="nb">set</span><span class="p">(</span><span class="n">d</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> |
| <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span> |
| <span class="sa">f</span><span class="s1">'All the dicts in the input data should have the same keys. '</span> |
| <span class="sa">f</span><span class="s1">'Got: </span><span class="si">{</span><span class="n">extra_keys</span><span class="si">}</span><span class="s1"> instead.'</span><span class="p">)</span> |
| <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">d</span><span class="o">.</span><span class="n">items</span><span class="p">():</span> |
| <span class="n">keys_to_element_list</span><span class="p">[</span><span class="n">key</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">keys_to_element_list</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_convert_dict_of_lists_to_lists_of_dict</span><span class="p">(</span> |
| <span class="n">dict_of_lists</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">]])</span> <span class="o">-></span> <span class="n">List</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]:</span> |
| <span class="n">batch_length</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="nb">next</span><span class="p">(</span><span class="nb">iter</span><span class="p">(</span><span class="n">dict_of_lists</span><span class="o">.</span><span class="n">values</span><span class="p">())))</span> |
| <span class="n">result</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="p">[{}</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">batch_length</span><span class="p">)]</span> |
| <span class="c1"># all the values in the dict_of_lists should have same length</span> |
| <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">values</span> <span class="ow">in</span> <span class="n">dict_of_lists</span><span class="o">.</span><span class="n">items</span><span class="p">():</span> |
| <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">values</span><span class="p">)</span> <span class="o">==</span> <span class="n">batch_length</span><span class="p">,</span> <span class="p">(</span> |
| <span class="s2">"This function expects all the values "</span> |
| <span class="s2">"in the dict_of_lists to have same length."</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">values</span><span class="p">)):</span> |
| <span class="n">result</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">values</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> |
| <span class="k">return</span> <span class="n">result</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_map_errors_to_beam_row</span><span class="p">(</span><span class="n">element</span><span class="p">,</span> <span class="n">cls_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span> |
| <span class="n">row_elements</span> <span class="o">=</span> <span class="p">{</span> |
| <span class="s1">'element'</span><span class="p">:</span> <span class="n">element</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> |
| <span class="s1">'msg'</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">element</span><span class="p">[</span><span class="mi">1</span><span class="p">][</span><span class="mi">1</span><span class="p">]),</span> |
| <span class="s1">'stack'</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">element</span><span class="p">[</span><span class="mi">1</span><span class="p">][</span><span class="mi">2</span><span class="p">]),</span> |
| <span class="p">}</span> |
| <span class="k">if</span> <span class="n">cls_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span> |
| <span class="n">row_elements</span><span class="p">[</span><span class="s1">'transform_name'</span><span class="p">]</span> <span class="o">=</span> <span class="n">cls_name</span> |
| <span class="k">return</span> <span class="n">beam</span><span class="o">.</span><span class="n">Row</span><span class="p">(</span><span class="o">**</span><span class="n">row_elements</span><span class="p">)</span> |
| |
| |
| <span class="k">class</span> <span class="nc">ArtifactMode</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span> |
| <span class="n">PRODUCE</span> <span class="o">=</span> <span class="s1">'produce'</span> |
| <span class="n">CONSUME</span> <span class="o">=</span> <span class="s1">'consume'</span> |
| |
| |
| <div class="viewcode-block" id="MLTransformProvider"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.MLTransformProvider">[docs]</a><span class="k">class</span> <span class="nc">MLTransformProvider</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Data processing transforms that are intended to be used with MLTransform</span> |
| <span class="sd"> should subclass MLTransformProvider and implement</span> |
| <span class="sd"> get_ptransform_for_processing().</span> |
| |
| <span class="sd"> get_ptransform_for_processing() method should return a PTransform that can be</span> |
| <span class="sd"> used to process the data.</span> |
| |
| <span class="sd"> """</span> |
| <div class="viewcode-block" id="MLTransformProvider.get_ptransform_for_processing"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.MLTransformProvider.get_ptransform_for_processing">[docs]</a> <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span> |
| <span class="k">def</span> <span class="nf">get_ptransform_for_processing</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span> <span class="o">-></span> <span class="n">beam</span><span class="o">.</span><span class="n">PTransform</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Returns a PTransform that can be used to process the data.</span> |
| <span class="sd"> """</span></div> |
| |
| <div class="viewcode-block" id="MLTransformProvider.get_counter"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.MLTransformProvider.get_counter">[docs]</a> <span class="k">def</span> <span class="nf">get_counter</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Returns the counter name for the data processing transform.</span> |
| <span class="sd"> """</span> |
| <span class="n">counter_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span> |
| <span class="k">return</span> <span class="n">Metrics</span><span class="o">.</span><span class="n">counter</span><span class="p">(</span><span class="n">MLTransform</span><span class="p">,</span> <span class="sa">f</span><span class="s1">'BeamML_</span><span class="si">{</span><span class="n">counter_name</span><span class="si">}</span><span class="s1">'</span><span class="p">)</span></div></div> |
| |
| |
| <div class="viewcode-block" id="BaseOperation"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.BaseOperation">[docs]</a><span class="k">class</span> <span class="nc">BaseOperation</span><span class="p">(</span><span class="n">Generic</span><span class="p">[</span><span class="n">OperationInputT</span><span class="p">,</span> <span class="n">OperationOutputT</span><span class="p">],</span> |
| <span class="n">MLTransformProvider</span><span class="p">,</span> |
| <span class="n">abc</span><span class="o">.</span><span class="n">ABC</span><span class="p">):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">columns</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">])</span> <span class="o">-></span> <span class="kc">None</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Base Opertation class data processing transformations.</span> |
| <span class="sd"> Args:</span> |
| <span class="sd"> columns: List of column names to apply the transformation.</span> |
| <span class="sd"> """</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">columns</span> <span class="o">=</span> <span class="n">columns</span> |
| |
| <div class="viewcode-block" id="BaseOperation.apply_transform"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.BaseOperation.apply_transform">[docs]</a> <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span> |
| <span class="k">def</span> <span class="nf">apply_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">:</span> <span class="n">OperationInputT</span><span class="p">,</span> |
| <span class="n">output_column_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-></span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">OperationOutputT</span><span class="p">]:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Define any processing logic in the apply_transform() method.</span> |
| <span class="sd"> processing logics are applied on inputs and returns a transformed</span> |
| <span class="sd"> output.</span> |
| <span class="sd"> Args:</span> |
| <span class="sd"> inputs: input data.</span> |
| <span class="sd"> """</span></div> |
| |
| <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">:</span> <span class="n">OperationInputT</span><span class="p">,</span> |
| <span class="n">output_column_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-></span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">OperationOutputT</span><span class="p">]:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> This method is called when the instance of the class is called.</span> |
| <span class="sd"> This method will invoke the apply() method of the class.</span> |
| <span class="sd"> """</span> |
| <span class="n">transformed_data</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">apply_transform</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">output_column_name</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">transformed_data</span></div> |
| |
| |
| <div class="viewcode-block" id="ProcessHandler"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.ProcessHandler">[docs]</a><span class="k">class</span> <span class="nc">ProcessHandler</span><span class="p">(</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">PTransform</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">],</span> |
| <span class="n">Union</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">MLTransformOutputT</span><span class="p">],</span> |
| <span class="n">Tuple</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">MLTransformOutputT</span><span class="p">],</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">Row</span><span class="p">]]]],</span> |
| <span class="n">abc</span><span class="o">.</span><span class="n">ABC</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Only for internal use. No backwards compatibility guarantees.</span> |
| <span class="sd"> """</span> |
| <div class="viewcode-block" id="ProcessHandler.append_transform"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.ProcessHandler.append_transform">[docs]</a> <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span> |
| <span class="k">def</span> <span class="nf">append_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">transform</span><span class="p">:</span> <span class="n">BaseOperation</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Append transforms to the ProcessHandler.</span> |
| <span class="sd"> """</span></div></div> |
| |
| |
| <span class="c1"># TODO:https://github.com/apache/beam/issues/29356</span> |
| <span class="c1"># Add support for inference_fn</span> |
| <div class="viewcode-block" id="EmbeddingsManager"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.EmbeddingsManager">[docs]</a><span class="k">class</span> <span class="nc">EmbeddingsManager</span><span class="p">(</span><span class="n">MLTransformProvider</span><span class="p">):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">columns</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> |
| <span class="o">*</span><span class="p">,</span> |
| <span class="c1"># common args for all ModelHandlers.</span> |
| <span class="n">load_model_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="n">min_batch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="n">max_batch_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="n">large_model</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> |
| <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">load_model_args</span> <span class="o">=</span> <span class="n">load_model_args</span> <span class="ow">or</span> <span class="p">{}</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">min_batch_size</span> <span class="o">=</span> <span class="n">min_batch_size</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">max_batch_size</span> <span class="o">=</span> <span class="n">max_batch_size</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">large_model</span> <span class="o">=</span> <span class="n">large_model</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">columns</span> <span class="o">=</span> <span class="n">columns</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">inference_args</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s1">'inference_args'</span><span class="p">,</span> <span class="p">{})</span> |
| |
| <span class="k">if</span> <span class="n">kwargs</span><span class="p">:</span> |
| <span class="n">_LOGGER</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">"Ignoring the following arguments: </span><span class="si">%s</span><span class="s2">"</span><span class="p">,</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> |
| |
| <span class="c1"># TODO:https://github.com/apache/beam/pull/29564 add set_model_handler method</span> |
| <div class="viewcode-block" id="EmbeddingsManager.get_model_handler"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.EmbeddingsManager.get_model_handler">[docs]</a> <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span> |
| <span class="k">def</span> <span class="nf">get_model_handler</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">ModelHandler</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Return framework specific model handler.</span> |
| <span class="sd"> """</span></div> |
| |
| <div class="viewcode-block" id="EmbeddingsManager.get_columns_to_apply"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.EmbeddingsManager.get_columns_to_apply">[docs]</a> <span class="k">def</span> <span class="nf">get_columns_to_apply</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">columns</span></div></div> |
| |
| |
| <div class="viewcode-block" id="MLTransform"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.MLTransform">[docs]</a><span class="k">class</span> <span class="nc">MLTransform</span><span class="p">(</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">PTransform</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">],</span> |
| <span class="n">Union</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">MLTransformOutputT</span><span class="p">],</span> |
| <span class="n">Tuple</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">MLTransformOutputT</span><span class="p">],</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">Row</span><span class="p">]]]],</span> |
| <span class="n">Generic</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">,</span> <span class="n">MLTransformOutputT</span><span class="p">]):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="o">*</span><span class="p">,</span> |
| <span class="n">write_artifact_location</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="n">read_artifact_location</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="n">transforms</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">MLTransformProvider</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> MLTransform is a Beam PTransform that can be used to apply</span> |
| <span class="sd"> transformations to the data. MLTransform is used to wrap the</span> |
| <span class="sd"> data processing transforms provided by Apache Beam. MLTransform</span> |
| <span class="sd"> works in two modes: write and read. In the write mode,</span> |
| <span class="sd"> MLTransform will apply the transforms to the data and store the</span> |
| <span class="sd"> artifacts in the write_artifact_location. In the read mode,</span> |
| <span class="sd"> MLTransform will read the artifacts from the</span> |
| <span class="sd"> read_artifact_location and apply the transforms to the data. The</span> |
| <span class="sd"> artifact location should be a valid storage path where the artifacts</span> |
| <span class="sd"> can be written to or read from.</span> |
| |
| <span class="sd"> Note that when consuming artifacts, it is not necessary to pass the</span> |
| <span class="sd"> transforms since they are inherently stored within the artifacts</span> |
| <span class="sd"> themselves.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> write_artifact_location: A storage location for artifacts resulting from</span> |
| <span class="sd"> MLTransform. These artifacts include transformations applied to</span> |
| <span class="sd"> the dataset and generated values like min, max from ScaleTo01,</span> |
| <span class="sd"> and mean, var from ScaleToZScore. Artifacts are produced and written</span> |
| <span class="sd"> to this location when using `write_artifact_mode`.</span> |
| <span class="sd"> Later MLTransforms can reuse produced artifacts by setting</span> |
| <span class="sd"> `read_artifact_mode` instead of `write_artifact_mode`. The value</span> |
| <span class="sd"> assigned to `write_artifact_location` should be a valid storage</span> |
| <span class="sd"> directory that the artifacts from this transform can be written to.</span> |
| <span class="sd"> If no directory exists at this location, one will be created. This will</span> |
| <span class="sd"> overwrite any artifacts already in this location, so distinct locations</span> |
| <span class="sd"> should be used for each instance of MLTransform. Only one of</span> |
| <span class="sd"> write_artifact_location and read_artifact_location should be specified.</span> |
| <span class="sd"> read_artifact_location: A storage location to read artifacts resulting</span> |
| <span class="sd"> froma previous MLTransform. These artifacts include transformations</span> |
| <span class="sd"> applied to the dataset and generated values like min, max from</span> |
| <span class="sd"> ScaleTo01, and mean, var from ScaleToZScore. Note that when consuming</span> |
| <span class="sd"> artifacts, it is not necessary to pass the transforms since they are</span> |
| <span class="sd"> inherently stored within the artifacts themselves. The value assigned</span> |
| <span class="sd"> to `read_artifact_location` should be a valid storage path where the</span> |
| <span class="sd"> artifacts can be read from. Only one of write_artifact_location and</span> |
| <span class="sd"> read_artifact_location should be specified.</span> |
| <span class="sd"> transforms: A list of transforms to apply to the data. All the transforms</span> |
| <span class="sd"> are applied in the order they are specified. The input of the</span> |
| <span class="sd"> i-th transform is the output of the (i-1)-th transform. Multi-input</span> |
| <span class="sd"> transforms are not supported yet.</span> |
| <span class="sd"> """</span> |
| <span class="k">if</span> <span class="n">read_artifact_location</span> <span class="ow">and</span> <span class="n">write_artifact_location</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span> |
| <span class="s1">'Only one of read_artifact_location or write_artifact_location can '</span> |
| <span class="s1">'be specified to initialize MLTransform'</span><span class="p">)</span> |
| |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">read_artifact_location</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">write_artifact_location</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span> |
| <span class="s1">'Either a read_artifact_location or write_artifact_location must be '</span> |
| <span class="s1">'specified to initialize MLTransform'</span><span class="p">)</span> |
| |
| <span class="k">if</span> <span class="n">read_artifact_location</span><span class="p">:</span> |
| <span class="n">artifact_location</span> <span class="o">=</span> <span class="n">read_artifact_location</span> |
| <span class="n">artifact_mode</span> <span class="o">=</span> <span class="n">ArtifactMode</span><span class="o">.</span><span class="n">CONSUME</span> |
| <span class="k">if</span> <span class="n">transforms</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span> |
| <span class="s1">'Transforms should not be passed in read mode. In read mode, '</span> |
| <span class="s1">'the transforms are read from the artifact location.'</span><span class="p">)</span> |
| |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">artifact_location</span> <span class="o">=</span> <span class="n">write_artifact_location</span> <span class="c1"># type: ignore[assignment]</span> |
| <span class="n">artifact_mode</span> <span class="o">=</span> <span class="n">ArtifactMode</span><span class="o">.</span><span class="n">PRODUCE</span> |
| |
| <span class="bp">self</span><span class="o">.</span><span class="n">_parent_artifact_location</span> <span class="o">=</span> <span class="n">artifact_location</span> |
| |
| <span class="bp">self</span><span class="o">.</span><span class="n">_artifact_mode</span> <span class="o">=</span> <span class="n">artifact_mode</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span> <span class="o">=</span> <span class="n">transforms</span> <span class="ow">or</span> <span class="p">[]</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_counter</span> <span class="o">=</span> <span class="n">Metrics</span><span class="o">.</span><span class="n">counter</span><span class="p">(</span> |
| <span class="n">MLTransform</span><span class="p">,</span> <span class="sa">f</span><span class="s1">'BeamML_</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s1">'</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_with_exception_handling</span> <span class="o">=</span> <span class="kc">False</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_exception_handling_args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]</span> <span class="o">=</span> <span class="p">{}</span> |
| |
| <div class="viewcode-block" id="MLTransform.expand"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.MLTransform.expand">[docs]</a> <span class="k">def</span> <span class="nf">expand</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> <span class="n">pcoll</span><span class="p">:</span> <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">ExampleT</span><span class="p">]</span> |
| <span class="p">)</span> <span class="o">-></span> <span class="n">Union</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">MLTransformOutputT</span><span class="p">],</span> |
| <span class="n">Tuple</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">MLTransformOutputT</span><span class="p">],</span> |
| <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">[</span><span class="n">beam</span><span class="o">.</span><span class="n">Row</span><span class="p">]]]:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> This is the entrypoint for the MLTransform. This method will</span> |
| <span class="sd"> invoke the process_data() method of the ProcessHandler instance</span> |
| <span class="sd"> to process the incoming data.</span> |
| |
| <span class="sd"> process_data takes in a PCollection and applies the PTransforms</span> |
| <span class="sd"> necessary to process the data and returns a PCollection of</span> |
| <span class="sd"> transformed data.</span> |
| <span class="sd"> Args:</span> |
| <span class="sd"> pcoll: A PCollection of ExampleT type.</span> |
| <span class="sd"> Returns:</span> |
| <span class="sd"> A PCollection of MLTransformOutputT type</span> |
| <span class="sd"> """</span> |
| <span class="n">upstream_errors</span> <span class="o">=</span> <span class="p">[]</span> |
| <span class="n">_</span> <span class="o">=</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">_validate_transform</span><span class="p">(</span><span class="n">transform</span><span class="p">)</span> <span class="k">for</span> <span class="n">transform</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span><span class="p">]</span> |
| <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_artifact_mode</span> <span class="o">==</span> <span class="n">ArtifactMode</span><span class="o">.</span><span class="n">PRODUCE</span><span class="p">:</span> |
| <span class="n">ptransform_partitioner</span> <span class="o">=</span> <span class="n">_MLTransformToPTransformMapper</span><span class="p">(</span> |
| <span class="n">transforms</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">transforms</span><span class="p">,</span> |
| <span class="n">artifact_location</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_parent_artifact_location</span><span class="p">,</span> |
| <span class="n">artifact_mode</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_artifact_mode</span><span class="p">,</span> |
| <span class="n">pipeline_options</span><span class="o">=</span><span class="n">pcoll</span><span class="o">.</span><span class="n">pipeline</span><span class="o">.</span><span class="n">options</span><span class="p">)</span> |
| <span class="n">ptransform_list</span> <span class="o">=</span> <span class="n">ptransform_partitioner</span><span class="o">.</span><span class="n">create_and_save_ptransform_list</span><span class="p">()</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">ptransform_list</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">_MLTransformToPTransformMapper</span><span class="o">.</span><span class="n">load_transforms_from_artifact_location</span><span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_parent_artifact_location</span><span class="p">))</span> |
| |
| <span class="c1"># the saved transforms has artifact mode set to PRODUCE.</span> |
| <span class="c1"># set the artifact mode to CONSUME.</span> |
| <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ptransform_list</span><span class="p">)):</span> |
| <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">ptransform_list</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="s1">'artifact_mode'</span><span class="p">):</span> |
| <span class="n">ptransform_list</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">artifact_mode</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_artifact_mode</span> |
| |
| <span class="n">transform_name</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="k">for</span> <span class="n">ptransform</span> <span class="ow">in</span> <span class="n">ptransform_list</span><span class="p">:</span> |
| <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_with_exception_handling</span><span class="p">:</span> |
| <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">ptransform</span><span class="p">,</span> <span class="s1">'with_exception_handling'</span><span class="p">):</span> |
| <span class="n">ptransform</span> <span class="o">=</span> <span class="n">ptransform</span><span class="o">.</span><span class="n">with_exception_handling</span><span class="p">(</span> |
| <span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">_exception_handling_args</span><span class="p">)</span> |
| <span class="n">pcoll</span><span class="p">,</span> <span class="n">bad_results</span> <span class="o">=</span> <span class="n">pcoll</span> <span class="o">|</span> <span class="n">ptransform</span> |
| <span class="c1"># RunInference outputs a RunInferenceDLQ instead of a PCollection.</span> |
| <span class="c1"># since TFTProcessHandler and RunInferene are supported, try to infer</span> |
| <span class="c1"># the type of bad_results and append it to the list of errors.</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">bad_results</span><span class="p">,</span> <span class="n">RunInferenceDLQ</span><span class="p">):</span> |
| <span class="n">bad_results</span> <span class="o">=</span> <span class="n">bad_results</span><span class="o">.</span><span class="n">failed_inferences</span> |
| <span class="n">transform_name</span> <span class="o">=</span> <span class="n">ptransform</span><span class="o">.</span><span class="n">annotations</span><span class="p">()[</span><span class="s1">'model_handler'</span><span class="p">]</span> |
| <span class="k">elif</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">bad_results</span><span class="p">,</span> <span class="n">beam</span><span class="o">.</span><span class="n">PCollection</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span> |
| <span class="sa">f</span><span class="s1">'Unexpected type for bad_results: </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">bad_results</span><span class="p">)</span><span class="si">}</span><span class="s1">'</span><span class="p">)</span> |
| <span class="n">bad_results</span> <span class="o">=</span> <span class="n">bad_results</span> <span class="o">|</span> <span class="n">beam</span><span class="o">.</span><span class="n">Map</span><span class="p">(</span> |
| <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">_map_errors_to_beam_row</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">transform_name</span><span class="p">))</span> |
| <span class="n">upstream_errors</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">bad_results</span><span class="p">)</span> |
| |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">pcoll</span> <span class="o">=</span> <span class="n">pcoll</span> <span class="o">|</span> <span class="n">ptransform</span> |
| <span class="n">_</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">pcoll</span><span class="o">.</span><span class="n">pipeline</span> |
| <span class="o">|</span> <span class="s2">"MLTransformMetricsUsage"</span> <span class="o">>></span> <span class="n">MLTransformMetricsUsage</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span> |
| |
| <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_with_exception_handling</span><span class="p">:</span> |
| <span class="n">bad_pcoll</span> <span class="o">=</span> <span class="p">(</span><span class="n">upstream_errors</span> <span class="o">|</span> <span class="n">beam</span><span class="o">.</span><span class="n">Flatten</span><span class="p">())</span> |
| <span class="k">return</span> <span class="n">pcoll</span><span class="p">,</span> <span class="n">bad_pcoll</span> <span class="c1"># type: ignore[return-value]</span> |
| <span class="k">return</span> <span class="n">pcoll</span> <span class="c1"># type: ignore[return-value]</span></div> |
| |
| <div class="viewcode-block" id="MLTransform.with_transform"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.MLTransform.with_transform">[docs]</a> <span class="k">def</span> <span class="nf">with_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">transform</span><span class="p">:</span> <span class="n">MLTransformProvider</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Add a transform to the MLTransform pipeline.</span> |
| <span class="sd"> Args:</span> |
| <span class="sd"> transform: A BaseOperation instance.</span> |
| <span class="sd"> Returns:</span> |
| <span class="sd"> A MLTransform instance.</span> |
| <span class="sd"> """</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_validate_transform</span><span class="p">(</span><span class="n">transform</span><span class="p">)</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">transform</span><span class="p">)</span> |
| <span class="k">return</span> <span class="bp">self</span></div> |
| |
| <span class="k">def</span> <span class="nf">_validate_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">transform</span><span class="p">):</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span> <span class="n">MLTransformProvider</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span> |
| <span class="s1">'transform must be a subclass of BaseOperation. '</span> |
| <span class="s1">'Got: </span><span class="si">%s</span><span class="s1"> instead.'</span> <span class="o">%</span> <span class="nb">type</span><span class="p">(</span><span class="n">transform</span><span class="p">))</span> |
| |
| <div class="viewcode-block" id="MLTransform.with_exception_handling"><a class="viewcode-back" href="../../../../apache_beam.ml.transforms.base.html#apache_beam.ml.transforms.base.MLTransform.with_exception_handling">[docs]</a> <span class="k">def</span> <span class="nf">with_exception_handling</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="n">exc_class</span><span class="o">=</span><span class="ne">Exception</span><span class="p">,</span> <span class="n">use_subprocess</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">threshold</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_with_exception_handling</span> <span class="o">=</span> <span class="kc">True</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_exception_handling_args</span> <span class="o">=</span> <span class="p">{</span> |
| <span class="s1">'exc_class'</span><span class="p">:</span> <span class="n">exc_class</span><span class="p">,</span> |
| <span class="s1">'use_subprocess'</span><span class="p">:</span> <span class="n">use_subprocess</span><span class="p">,</span> |
| <span class="s1">'threshold'</span><span class="p">:</span> <span class="n">threshold</span> |
| <span class="p">}</span> |
| <span class="k">return</span> <span class="bp">self</span></div></div> |
| |
| |
| <span class="k">class</span> <span class="nc">MLTransformMetricsUsage</span><span class="p">(</span><span class="n">beam</span><span class="o">.</span><span class="n">PTransform</span><span class="p">):</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ml_transform</span><span class="p">:</span> <span class="n">MLTransform</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_ml_transform</span> <span class="o">=</span> <span class="n">ml_transform</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_ml_transform</span><span class="o">.</span><span class="n">_counter</span><span class="o">.</span><span class="n">inc</span><span class="p">()</span> |
| |
| <span class="k">def</span> <span class="nf">expand</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pipeline</span><span class="p">):</span> |
| <span class="k">def</span> <span class="nf">_increment_counters</span><span class="p">():</span> |
| <span class="c1"># increment for MLTransform.</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_ml_transform</span><span class="o">.</span><span class="n">_counter</span><span class="o">.</span><span class="n">inc</span><span class="p">()</span> |
| <span class="c1"># increment if data processing transforms are passed.</span> |
| <span class="n">transforms</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_ml_transform</span><span class="o">.</span><span class="n">transforms</span> |
| <span class="k">if</span> <span class="n">transforms</span><span class="p">:</span> |
| <span class="k">for</span> <span class="n">transform</span> <span class="ow">in</span> <span class="n">transforms</span><span class="p">:</span> |
| <span class="n">transform</span><span class="o">.</span><span class="n">get_counter</span><span class="p">()</span><span class="o">.</span><span class="n">inc</span><span class="p">()</span> |
| |
| <span class="n">_</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">pipeline</span> |
| <span class="o">|</span> <span class="n">beam</span><span class="o">.</span><span class="n">Create</span><span class="p">([</span><span class="kc">None</span><span class="p">])</span> |
| <span class="o">|</span> <span class="n">beam</span><span class="o">.</span><span class="n">Map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">_</span><span class="p">:</span> <span class="n">_increment_counters</span><span class="p">()))</span> |
| |
| |
| <span class="k">class</span> <span class="nc">_TransformAttributeManager</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Base class used for saving and loading the attributes.</span> |
| <span class="sd"> """</span> |
| <span class="nd">@staticmethod</span> |
| <span class="k">def</span> <span class="nf">save_attributes</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Save the attributes to json file using stdlib json.</span> |
| <span class="sd"> """</span> |
| <span class="k">raise</span> <span class="ne">NotImplementedError</span> |
| |
| <span class="nd">@staticmethod</span> |
| <span class="k">def</span> <span class="nf">load_attributes</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Load the attributes from json file.</span> |
| <span class="sd"> """</span> |
| <span class="k">raise</span> <span class="ne">NotImplementedError</span> |
| |
| |
| <span class="k">class</span> <span class="nc">_JsonPickleTransformAttributeManager</span><span class="p">(</span><span class="n">_TransformAttributeManager</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Use Jsonpickle to save and load the attributes. Here the attributes refer</span> |
| <span class="sd"> to the list of PTransforms that are used to process the data.</span> |
| |
| <span class="sd"> jsonpickle is used to serialize the PTransforms and save it to a json file and</span> |
| <span class="sd"> is compatible across python versions.</span> |
| <span class="sd"> """</span> |
| <span class="nd">@staticmethod</span> |
| <span class="k">def</span> <span class="nf">_is_remote_path</span><span class="p">(</span><span class="n">path</span><span class="p">):</span> |
| <span class="n">is_gcs</span> <span class="o">=</span> <span class="n">path</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s1">'gs://'</span><span class="p">)</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span> |
| <span class="c1"># TODO:https://github.com/apache/beam/issues/29356</span> |
| <span class="c1"># Add support for other remote paths.</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">is_gcs</span> <span class="ow">and</span> <span class="n">path</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s1">'://'</span><span class="p">)</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span> |
| <span class="s2">"Artifact locations are currently supported for only available for "</span> |
| <span class="s2">"local paths and GCS paths. Got: </span><span class="si">%s</span><span class="s2">"</span> <span class="o">%</span> <span class="n">path</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">is_gcs</span> |
| |
| <span class="nd">@staticmethod</span> |
| <span class="k">def</span> <span class="nf">save_attributes</span><span class="p">(</span> |
| <span class="n">ptransform_list</span><span class="p">,</span> |
| <span class="n">artifact_location</span><span class="p">,</span> |
| <span class="o">**</span><span class="n">kwargs</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="c1"># if an artifact location is present, instead of overwriting the</span> |
| <span class="c1"># existing file, raise an error since the same artifact location</span> |
| <span class="c1"># can be used by multiple beam jobs and this could result in undesired</span> |
| <span class="c1"># behavior.</span> |
| <span class="k">if</span> <span class="n">FileSystems</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">FileSystems</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">,</span> |
| <span class="n">_ATTRIBUTE_FILE_NAME</span><span class="p">)):</span> |
| <span class="k">raise</span> <span class="ne">FileExistsError</span><span class="p">(</span> |
| <span class="s2">"The artifact location </span><span class="si">%s</span><span class="s2"> already exists and contains </span><span class="si">%s</span><span class="s2">. Please "</span> |
| <span class="s2">"specify a different location."</span> <span class="o">%</span> |
| <span class="p">(</span><span class="n">artifact_location</span><span class="p">,</span> <span class="n">_ATTRIBUTE_FILE_NAME</span><span class="p">))</span> |
| |
| <span class="k">if</span> <span class="n">_JsonPickleTransformAttributeManager</span><span class="o">.</span><span class="n">_is_remote_path</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">):</span> |
| <span class="n">temp_dir</span> <span class="o">=</span> <span class="n">tempfile</span><span class="o">.</span><span class="n">mkdtemp</span><span class="p">()</span> |
| <span class="n">temp_json_file</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">temp_dir</span><span class="p">,</span> <span class="n">_ATTRIBUTE_FILE_NAME</span><span class="p">)</span> |
| <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">temp_json_file</span><span class="p">,</span> <span class="s1">'w+'</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span> |
| <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">jsonpickle</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">ptransform_list</span><span class="p">))</span> |
| <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">temp_json_file</span><span class="p">,</span> <span class="s1">'rb'</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span> |
| <span class="kn">from</span> <span class="nn">apache_beam.runners.dataflow.internal</span> <span class="kn">import</span> <span class="n">apiclient</span> |
| <span class="n">_LOGGER</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s1">'Creating artifact location: </span><span class="si">%s</span><span class="s1">'</span><span class="p">,</span> <span class="n">artifact_location</span><span class="p">)</span> |
| <span class="c1"># pipeline options required to for the client to configure project.</span> |
| <span class="n">options</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">'options'</span><span class="p">)</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="n">apiclient</span><span class="o">.</span><span class="n">DataflowApplicationClient</span><span class="p">(</span><span class="n">options</span><span class="o">=</span><span class="n">options</span><span class="p">)</span><span class="o">.</span><span class="n">stage_file</span><span class="p">(</span> |
| <span class="n">gcs_or_local_path</span><span class="o">=</span><span class="n">artifact_location</span><span class="p">,</span> |
| <span class="n">file_name</span><span class="o">=</span><span class="n">_ATTRIBUTE_FILE_NAME</span><span class="p">,</span> |
| <span class="n">stream</span><span class="o">=</span><span class="n">f</span><span class="p">,</span> |
| <span class="n">mime_type</span><span class="o">=</span><span class="s1">'application/json'</span><span class="p">)</span> |
| <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">exc</span><span class="p">:</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">options</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span> |
| <span class="s2">"Failed to create Dataflow client. "</span> |
| <span class="s2">"Pipeline options are required to save the attributes."</span> |
| <span class="s2">"in the artifact location </span><span class="si">%s</span><span class="s2">"</span> <span class="o">%</span> <span class="n">artifact_location</span><span class="p">)</span> <span class="kn">from</span> <span class="nn">exc</span> |
| <span class="k">raise</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">FileSystems</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">):</span> |
| <span class="n">FileSystems</span><span class="o">.</span><span class="n">mkdirs</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">)</span> |
| <span class="c1"># FileSystems.open() fails if the file does not exist.</span> |
| <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">,</span> <span class="n">_ATTRIBUTE_FILE_NAME</span><span class="p">),</span> |
| <span class="s1">'w+'</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span> |
| <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">jsonpickle</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="n">ptransform_list</span><span class="p">))</span> |
| |
| <span class="nd">@staticmethod</span> |
| <span class="k">def</span> <span class="nf">load_attributes</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">):</span> |
| <span class="k">with</span> <span class="n">FileSystems</span><span class="o">.</span><span class="n">open</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">,</span> <span class="n">_ATTRIBUTE_FILE_NAME</span><span class="p">),</span> |
| <span class="s1">'rb'</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span> |
| <span class="k">return</span> <span class="n">jsonpickle</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="n">f</span><span class="o">.</span><span class="n">read</span><span class="p">())</span> |
| |
| |
| <span class="n">_transform_attribute_manager</span> <span class="o">=</span> <span class="n">_JsonPickleTransformAttributeManager</span> |
| |
| |
| <span class="k">class</span> <span class="nc">_MLTransformToPTransformMapper</span><span class="p">:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> This class takes in a list of data processing transforms compatible to be</span> |
| <span class="sd"> wrapped around MLTransform and returns a list of PTransforms that are used to</span> |
| <span class="sd"> run the data processing transforms.</span> |
| |
| <span class="sd"> The _MLTransformToPTransformMapper is responsible for loading and saving the</span> |
| <span class="sd"> PTransforms or attributes of PTransforms to the artifact location to seal</span> |
| <span class="sd"> the gap between the training and inference pipelines.</span> |
| <span class="sd"> """</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">transforms</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">MLTransformProvider</span><span class="p">],</span> |
| <span class="n">artifact_location</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> |
| <span class="n">artifact_mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="n">ArtifactMode</span><span class="o">.</span><span class="n">PRODUCE</span><span class="p">,</span> |
| <span class="n">pipeline_options</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">PipelineOptions</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span> <span class="o">=</span> <span class="n">transforms</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_parent_artifact_location</span> <span class="o">=</span> <span class="n">artifact_location</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">artifact_mode</span> <span class="o">=</span> <span class="n">artifact_mode</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">pipeline_options</span> <span class="o">=</span> <span class="n">pipeline_options</span> |
| |
| <span class="k">def</span> <span class="nf">create_and_save_ptransform_list</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="n">ptransform_list</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">create_ptransform_list</span><span class="p">()</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">save_transforms_in_artifact_location</span><span class="p">(</span><span class="n">ptransform_list</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">ptransform_list</span> |
| |
| <span class="k">def</span> <span class="nf">create_ptransform_list</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="n">previous_ptransform_type</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="n">current_ptransform</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="n">ptransform_list</span> <span class="o">=</span> <span class="p">[]</span> |
| <span class="k">for</span> <span class="n">transform</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span><span class="p">:</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span> <span class="n">MLTransformProvider</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span> |
| <span class="s1">'Transforms must be instances of MLTransformProvider and '</span> |
| <span class="s1">'implement get_ptransform_for_processing() method.'</span><span class="p">)</span> |
| <span class="c1"># for each instance of PTransform, create a new artifact location</span> |
| <span class="n">current_ptransform</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">get_ptransform_for_processing</span><span class="p">(</span> |
| <span class="n">artifact_location</span><span class="o">=</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_parent_artifact_location</span><span class="p">,</span> <span class="n">uuid</span><span class="o">.</span><span class="n">uuid4</span><span class="p">()</span><span class="o">.</span><span class="n">hex</span><span class="p">[:</span><span class="mi">6</span><span class="p">]),</span> |
| <span class="n">artifact_mode</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">artifact_mode</span><span class="p">)</span> |
| <span class="n">append_transform</span> <span class="o">=</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">current_ptransform</span><span class="p">,</span> <span class="s1">'append_transform'</span><span class="p">)</span> |
| <span class="k">if</span> <span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">current_ptransform</span><span class="p">)</span> <span class="o">!=</span> |
| <span class="n">previous_ptransform_type</span><span class="p">)</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">append_transform</span><span class="p">:</span> |
| <span class="n">ptransform_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">current_ptransform</span><span class="p">)</span> |
| <span class="n">previous_ptransform_type</span> <span class="o">=</span> <span class="nb">type</span><span class="p">(</span><span class="n">current_ptransform</span><span class="p">)</span> |
| <span class="c1"># If different PTransform is appended to the list and the PTransform</span> |
| <span class="c1"># supports append_transform, append the transform to the PTransform.</span> |
| <span class="k">if</span> <span class="n">append_transform</span><span class="p">:</span> |
| <span class="n">ptransform_list</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">append_transform</span><span class="p">(</span><span class="n">transform</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">ptransform_list</span> |
| |
| <span class="k">def</span> <span class="nf">save_transforms_in_artifact_location</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ptransform_list</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Save the ptransform references to json file.</span> |
| <span class="sd"> """</span> |
| <span class="n">_transform_attribute_manager</span><span class="o">.</span><span class="n">save_attributes</span><span class="p">(</span> |
| <span class="n">ptransform_list</span><span class="o">=</span><span class="n">ptransform_list</span><span class="p">,</span> |
| <span class="n">artifact_location</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_parent_artifact_location</span><span class="p">,</span> |
| <span class="n">options</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">pipeline_options</span><span class="p">)</span> |
| |
| <span class="nd">@staticmethod</span> |
| <span class="k">def</span> <span class="nf">load_transforms_from_artifact_location</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">):</span> |
| <span class="k">return</span> <span class="n">_transform_attribute_manager</span><span class="o">.</span><span class="n">load_attributes</span><span class="p">(</span><span class="n">artifact_location</span><span class="p">)</span> |
| |
| |
| <span class="k">class</span> <span class="nc">_TextEmbeddingHandler</span><span class="p">(</span><span class="n">ModelHandler</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> A ModelHandler intended to be work on list[dict[str, str]] inputs.</span> |
| |
| <span class="sd"> The inputs to the model handler are expected to be a list of dicts.</span> |
| |
| <span class="sd"> For example, if the original mode is used with RunInference to take a</span> |
| <span class="sd"> PCollection[E] to a PCollection[P], this ModelHandler would take a</span> |
| <span class="sd"> PCollection[Dict[str, E]] to a PCollection[Dict[str, P]].</span> |
| |
| <span class="sd"> _TextEmbeddingHandler will accept an EmbeddingsManager instance, which</span> |
| <span class="sd"> contains the details of the model to be loaded and the inference_fn to be</span> |
| <span class="sd"> used. The purpose of _TextEmbeddingHandler is to generate embeddings for</span> |
| <span class="sd"> text inputs using the EmbeddingsManager instance.</span> |
| |
| <span class="sd"> If the input is not a text column, a RuntimeError will be raised.</span> |
| |
| <span class="sd"> This is an internal class and offers no backwards compatibility guarantees.</span> |
| |
| <span class="sd"> Args:</span> |
| <span class="sd"> embeddings_manager: An EmbeddingsManager instance.</span> |
| <span class="sd"> """</span> |
| <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">embeddings_manager</span><span class="p">:</span> <span class="n">EmbeddingsManager</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">embedding_config</span> <span class="o">=</span> <span class="n">embeddings_manager</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_underlying</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_config</span><span class="o">.</span><span class="n">get_model_handler</span><span class="p">()</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">columns</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_config</span><span class="o">.</span><span class="n">get_columns_to_apply</span><span class="p">()</span> |
| |
| <span class="k">def</span> <span class="nf">load_model</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="n">model</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_underlying</span><span class="o">.</span><span class="n">load_model</span><span class="p">()</span> |
| <span class="k">return</span> <span class="n">model</span> |
| |
| <span class="k">def</span> <span class="nf">_validate_column_data</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">):</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="nb">bytes</span><span class="p">)):</span> |
| <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span> |
| <span class="s1">'Embeddings can only be generated on Dict[str, str].'</span> |
| <span class="sa">f</span><span class="s1">'Got Dict[str, </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span><span class="si">}</span><span class="s1">] instead.'</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">_validate_batch</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]]):</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">batch</span> <span class="ow">or</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="nb">dict</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span> |
| <span class="s1">'Expected data to be dicts, got '</span> |
| <span class="sa">f</span><span class="s1">'</span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span><span class="si">}</span><span class="s1"> instead.'</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">_process_batch</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">dict_batch</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">]],</span> |
| <span class="n">model</span><span class="p">:</span> <span class="n">ModelT</span><span class="p">,</span> |
| <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]])</span> <span class="o">-></span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">]]:</span> |
| <span class="n">result</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="n">collections</span><span class="o">.</span><span class="n">defaultdict</span><span class="p">(</span><span class="nb">list</span><span class="p">)</span> |
| <span class="n">input_keys</span> <span class="o">=</span> <span class="n">dict_batch</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> |
| <span class="n">missing_columns_in_data</span> <span class="o">=</span> <span class="nb">set</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">columns</span><span class="p">)</span> <span class="o">-</span> <span class="nb">set</span><span class="p">(</span><span class="n">input_keys</span><span class="p">)</span> |
| <span class="k">if</span> <span class="n">missing_columns_in_data</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span> |
| <span class="sa">f</span><span class="s1">'Data does not contain the following columns '</span> |
| <span class="sa">f</span><span class="s1">': </span><span class="si">{</span><span class="n">missing_columns_in_data</span><span class="si">}</span><span class="s1">.'</span><span class="p">)</span> |
| <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">batch</span> <span class="ow">in</span> <span class="n">dict_batch</span><span class="o">.</span><span class="n">items</span><span class="p">():</span> |
| <span class="k">if</span> <span class="n">key</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">columns</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_validate_column_data</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> |
| <span class="n">prediction</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_underlying</span><span class="o">.</span><span class="n">run_inference</span><span class="p">(</span> |
| <span class="n">batch</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">)</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">prediction</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">):</span> |
| <span class="n">prediction</span> <span class="o">=</span> <span class="n">prediction</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span> |
| <span class="n">result</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">prediction</span> <span class="c1"># type: ignore[assignment]</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">result</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">prediction</span> <span class="c1"># type: ignore[assignment]</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">result</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">batch</span> |
| <span class="k">return</span> <span class="n">result</span> |
| |
| <span class="k">def</span> <span class="nf">run_inference</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">batch</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]],</span> |
| <span class="n">model</span><span class="p">:</span> <span class="n">ModelT</span><span class="p">,</span> |
| <span class="n">inference_args</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> |
| <span class="p">)</span> <span class="o">-></span> <span class="n">List</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Union</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]]]:</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Runs inference on a batch of text inputs. The inputs are expected to be</span> |
| <span class="sd"> a list of dicts. Each dict should have the same keys, and the shape</span> |
| <span class="sd"> should be of the same size for a single key across the batch.</span> |
| <span class="sd"> """</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_validate_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> |
| <span class="n">dict_batch</span> <span class="o">=</span> <span class="n">_convert_list_of_dicts_to_dict_of_lists</span><span class="p">(</span><span class="n">list_of_dicts</span><span class="o">=</span><span class="n">batch</span><span class="p">)</span> |
| <span class="n">transformed_batch</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_process_batch</span><span class="p">(</span><span class="n">dict_batch</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">inference_args</span><span class="p">)</span> |
| <span class="k">return</span> <span class="n">_convert_dict_of_lists_to_lists_of_dict</span><span class="p">(</span> |
| <span class="n">dict_of_lists</span><span class="o">=</span><span class="n">transformed_batch</span><span class="p">,</span> |
| <span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">get_metrics_namespace</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="nb">str</span><span class="p">:</span> |
| <span class="k">return</span> <span class="p">(</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_underlying</span><span class="o">.</span><span class="n">get_metrics_namespace</span><span class="p">()</span> <span class="ow">or</span> |
| <span class="s1">'BeamML_TextEmbeddingHandler'</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">batch_elements_kwargs</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-></span> <span class="n">Mapping</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]:</span> |
| <span class="n">batch_sizes_map</span> <span class="o">=</span> <span class="p">{}</span> |
| <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_config</span><span class="o">.</span><span class="n">max_batch_size</span><span class="p">:</span> |
| <span class="n">batch_sizes_map</span><span class="p">[</span><span class="s1">'max_batch_size'</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_config</span><span class="o">.</span><span class="n">max_batch_size</span> |
| <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_config</span><span class="o">.</span><span class="n">min_batch_size</span><span class="p">:</span> |
| <span class="n">batch_sizes_map</span><span class="p">[</span><span class="s1">'min_batch_size'</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">embedding_config</span><span class="o">.</span><span class="n">min_batch_size</span> |
| <span class="k">return</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_underlying</span><span class="o">.</span><span class="n">batch_elements_kwargs</span><span class="p">()</span> <span class="ow">or</span> <span class="n">batch_sizes_map</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_underlying</span><span class="o">.</span><span class="fm">__repr__</span><span class="p">()</span> |
| |
| <span class="k">def</span> <span class="nf">validate_inference_args</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">_</span><span class="p">):</span> |
| <span class="k">pass</span> |
| </pre></div> |
| |
| </div> |
| |
| </div> |
| <footer> |
| |
| |
| <hr/> |
| |
| <div role="contentinfo"> |
| <p> |
| © Copyright |
| |
| </p> |
| </div> |
| Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>. |
| |
| </footer> |
| |
| </div> |
| </div> |
| |
| </section> |
| |
| </div> |
| |
| |
| |
| <script type="text/javascript"> |
| jQuery(function () { |
| SphinxRtdTheme.Navigation.enable(true); |
| }); |
| </script> |
| |
| |
| |
| |
| |
| |
| </body> |
| </html> |
