nix-shell for tensorflow hub and google Bert

shell.nix

{ pkgs ? import <nixpkgs> {} }:
let
    bertTfHubModule = builtins.fetchTarball {
       url="https://tfhub.dev/google/bert_multi_cased_L-12_H-768_A-12/1?tf-hub-format=compressed";
       sha256="sha256:01svqqhs4jlpi0lax3kr4cj43kfjvvvzlhd2gkb61849hish3qk3";
    };

    googleBertSrc =  pkgs.fetchFromGitHub {
        owner = "google-research";
        repo = "bert";
        rev = "d66a146741588fb208450bde15aa7db143baaa69";
        sha256 = "1mdrr9xkfn3gzmbvpj4pwb2kll95pcfdb6q288xr3m9s0481radz";
      };

    packageOverrides = self: super: {
        tensorflow-hub = pkgs.python36Packages.buildPythonPackage rec {
          pname = "tensorflow-hub";
          version = "0.4.0";
          format = "wheel";

          src = pkgs.python36Packages.fetchPypi {
            pname = "tensorflow_hub";
            version = version;
            format = "wheel";
            sha256 = "299967392aeda7032d6bd29b178f65a9bd8e43cfc550ea427fd37507711d5ec2";
          };

          propagatedBuildInputs = [ self.tensorflow ];

          doCheck = false;  # too much
        };
      };

    python = pkgs.python36.override { inherit packageOverrides; };

    pythonWithPackages = python.withPackages (p: [ p.tensorflow-hub ]);

    bert = pkgs.writeScriptBin "run_bert" ''
      #! ${pkgs.stdenv.shell}
      ${pythonWithPackages.interpreter} ${googleBertSrc}/run_classifier_with_tfhub.py --bert_hub_module_handle="${bertTfHubModule}" "$@"
      '';
in
  pkgs.mkShell { buildInputs = [ bert ]; }