diff --git a/tensorflow/lite/tools/pip_package/Dockerfile b/tensorflow/lite/tools/pip_package/Dockerfile
index e4d47d0fa0d..3b92596758b 100644
--- a/tensorflow/lite/tools/pip_package/Dockerfile
+++ b/tensorflow/lite/tools/pip_package/Dockerfile
@@ -14,6 +14,7 @@ RUN apt-get update && \
       python-setuptools \
       python-wheel \
       python-numpy \
+      python-pip \
       libpython-dev \
       libpython-dev:armhf \
       libpython-dev:arm64 \
@@ -21,6 +22,7 @@ RUN apt-get update && \
       python3-setuptools \
       python3-wheel \
       python3-numpy \
+      python3-pip \
       libpython3-dev \
       libpython3-dev:armhf \
       libpython3-dev:arm64 \
@@ -29,8 +31,11 @@ RUN apt-get update && \
       zlib1g-dev  \
       zlib1g-dev:armhf \
       zlib1g-dev:arm64 \
-      swig \
       curl \
       unzip \
       git && \
     apt-get clean
+RUN pip install pip --upgrade
+RUN pip install pybind11
+RUN pip3 install pip --upgrade
+RUN pip3 install pybind11
diff --git a/tensorflow/lite/tools/pip_package/setup.py b/tensorflow/lite/tools/pip_package/setup.py
index f99a5b043dc..ebd4cfb3126 100644
--- a/tensorflow/lite/tools/pip_package/setup.py
+++ b/tensorflow/lite/tools/pip_package/setup.py
@@ -24,9 +24,11 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import glob
 import multiprocessing
 import os
 import subprocess
+import sys
 
 from distutils.command.build_ext import build_ext
 import numpy
@@ -65,7 +67,7 @@ for name in ['TARGET', 'TARGET_ARCH', 'CC_PREFIX', 'EXTRA_CXXFLAGS']:
 # with more than 4GB, use all the CPUs, otherwise only 1.
 def get_build_cpus():
   physical_bytes = os.sysconf('SC_PAGESIZE') * os.sysconf('SC_PHYS_PAGES')
-  if physical_bytes < (1<<30) * 4:
+  if physical_bytes < (1 << 30) * 4:
     return 1
   else:
     return multiprocessing.cpu_count()
@@ -73,9 +75,9 @@ def get_build_cpus():
 
 def make_args(target='', quiet=True):
   """Construct make command line."""
-  args = (['make', 'SHELL=/bin/bash',
-           'BUILD_WITH_NNAPI=false', '-C', TENSORFLOW_DIR]
-          + MAKE_CROSS_OPTIONS +
+  args = ([
+      'make', 'SHELL=/bin/bash', 'BUILD_WITH_NNAPI=false', '-C', TENSORFLOW_DIR
+  ] + MAKE_CROSS_OPTIONS +
           ['-f', RELATIVE_MAKEFILE_PATH, '-j',
            str(get_build_cpus())])
   if quiet:
@@ -128,28 +130,46 @@ class CustomBuildPy(build_py, object):
     return super(CustomBuildPy, self).run()
 
 
+def get_pybind_include():
+  """pybind11 include directory is not correctly resolved.
+
+  This fixes include directory to /usr/local/pythonX.X
+
+  Returns:
+    include directories to find pybind11
+  """
+  if sys.version_info[0] == 3:
+    include_dirs = glob.glob('/usr/local/include/python3*')
+  else:
+    include_dirs = glob.glob('/usr/local/include/python2*')
+  for include_dir in include_dirs:
+    os.symlink(include_dir, os.path.join(include_dir, 'include'))
+
+  return include_dirs
+
+
 LIB_TFLITE = 'tensorflow-lite'
 LIB_TFLITE_DIR = make_output('libdir')
 
 ext = Extension(
     name='%s._interpreter_wrapper' % PACKAGE_NAME,
     language='c++',
-    sources=['interpreter_wrapper/interpreter_wrapper.i',
-             'interpreter_wrapper/interpreter_wrapper.cc',
-             'interpreter_wrapper/numpy.cc',
-             'interpreter_wrapper/python_error_reporter.cc',
-             'interpreter_wrapper/python_utils.cc'],
+    sources=[
+        'interpreter_wrapper/interpreter_wrapper.cc',
+        'interpreter_wrapper/interpreter_wrapper_pybind11.cc',
+        'interpreter_wrapper/numpy.cc',
+        'interpreter_wrapper/python_error_reporter.cc',
+        'interpreter_wrapper/python_utils.cc'
+    ],
     extra_compile_args=['--std=c++11'],
-    swig_opts=['-c++',
-               '-I%s' % TENSORFLOW_DIR,
-               '-module', 'interpreter_wrapper',
-               '-outdir', PACKAGE_NAME],
-    include_dirs=[TENSORFLOW_DIR,
-                  os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
-                               'pip_package'),
-                  numpy.get_include(),
-                  os.path.join(DOWNLOADS_DIR, 'flatbuffers', 'include'),
-                  os.path.join(DOWNLOADS_DIR, 'absl')],
+    include_dirs=[
+        TENSORFLOW_DIR,
+        os.path.join(TENSORFLOW_DIR, 'tensorflow', 'lite', 'tools',
+                     'pip_package'),
+        numpy.get_include(),
+        os.path.join(DOWNLOADS_DIR, 'flatbuffers', 'include'),
+        os.path.join(DOWNLOADS_DIR, 'absl')
+    ] + get_pybind_include(),
     libraries=[LIB_TFLITE],
     library_dirs=[LIB_TFLITE_DIR])
 
@@ -186,9 +206,9 @@ setup(
     ext_modules=[ext],
     install_requires=[
         'numpy >= 1.16.0',
+        'pybind11 >= 2.4.3',
     ],
     cmdclass={
         'build_ext': CustomBuildExt,
         'build_py': CustomBuildPy,
-    }
-)
+    })