diff --git a/karma.conf.js b/karma.conf.js
index 2b69fef..b93f870 100644
--- a/karma.conf.js
+++ b/karma.conf.js
@@ -32,15 +32,27 @@ module.exports = function(config) {
 
     webpack: {
         module: {
-            loaders: [
-                {test: /\.(js|jsx)$/, loaders: ['babel']},
-            ]
+          rules: [
+            {
+              test: /\.(js|jsx)$/,
+              use: {
+                loader: 'babel-loader',
+                options: {
+                  presets: ['env'],
+                  plugins: [
+                    'transform-class-properties',
+                    'transform-react-jsx',
+                  ],
+                },
+              },
+            },
+          ],
         },
         externals: {
             react: 'React'
         },
         resolve: {
-          root: __dirname
+          modules: [__dirname, 'node_modules']
         }
     },
     webpackServer: {
diff --git a/package.json b/package.json
index 679988d..b72ca1d 100644
--- a/package.json
+++ b/package.json
@@ -5,7 +5,7 @@
   "main": "index.js",
   "scripts": {
     "start": "gulp server",
-    "prepublish": "babel ./src --out-dir ./lib",
+    "prepublish": "babel ./src --out-dir ./lib --plugins=transform-class-properties,transform-react-jsx --presets=env",
     "test": "karma start --single-run",
     "dev-test": "karma start"
   },
@@ -27,16 +27,17 @@
   },
   "homepage": "https://github.com/akiran/react-highlight",
   "dependencies": {
-    "highlight.js": "^9.11.0",
-    "react": "^15.5.4",
-    "react-dom": "^15.5.4"
+    "highlight.js": "^9.11.0"
   },
   "devDependencies": {
     "autoprefixer": "^6.7.7",
-    "babel": "^5.8.38",
-    "babel-core": "^5.8.38",
+    "babel-cli": "^6.26.0",
+    "babel-core": "^6.26.0",
     "babel-eslint": "^6.1.2",
-    "babel-loader": "^5.4.2",
+    "babel-loader": "^7.1.2",
+    "babel-plugin-transform-class-properties": "^6.24.1",
+    "babel-plugin-transform-react-jsx": "^6.24.1",
+    "babel-preset-env": "^1.6.0",
     "es5-shim": "^4.5.9",
     "eslint": "^3.19.0",
     "eslint-plugin-react": "^6.10.3",
@@ -54,7 +55,13 @@
     "node-libs-browser": "^2.0.0",
     "phantomjs-prebuilt": "^2.1.14",
     "raw-loader": "^0.5.1",
-    "webpack": "^1.15.0",
-    "webpack-dev-server": "^1.16.3"
+    "react": "^15.5.4",
+    "react-dom": "^15.5.4",
+    "webpack": "^3.6.0",
+    "webpack-dev-server": "^2.9.1"
+  },
+  "peerDependencies": {
+    "react": "^15.0.0 || ^16.0.0",
+    "react-dom": "^15.0.0 || ^16.0.0"
   }
 }
diff --git a/src/index.js b/src/index.js
index cb178f6..a5fff66 100644
--- a/src/index.js
+++ b/src/index.js
@@ -1,6 +1,5 @@
 import hljs from 'highlight.js';
 import React from 'react';
-import ReactDOM from 'react-dom';
 
 class Highlight extends React.Component {
     componentDidMount() {
@@ -12,32 +11,33 @@ class Highlight extends React.Component {
     }
 
     highlightCode() {
-        const domNode = ReactDOM.findDOMNode(this);
-        const nodes = domNode.querySelectorAll('pre code');
+        const nodes = this.el.querySelectorAll('pre code');
 
-        let i;
-        for (i = 0; i < nodes.length; i++) {
-            hljs.highlightBlock(nodes[i]);
+        for (let i = 0; i < nodes.length; i++) {
+            hljs.highlightBlock(nodes[i])
         }
     }
 
+    setEl = (el) => {
+        this.el = el;
+    };
+
     render() {
-        const {children, className, element, innerHTML} = this.props;
-        let Element = element ? React.DOM[element] : null;
+        const {children, className, element: Element, innerHTML} = this.props;
+        const props = { ref: this.setEl, className };
 
         if (innerHTML) {
-            if (!Element) {
-                Element = React.DOM.div
-            }
-
-            return Element({dangerouslySetInnerHTML: {__html: children}, className: className || null}, null);
-        } else {
+            props.dangerouslySetInnerHTML = { __html: children };
             if (Element) {
-                return Element({className}, children);
-            } else {
-                return <pre><code className={className}>{children}</code></pre>;
+                return <Element {...props} />;
             }
+            return <div {...props} />;
+        }
+
+        if (Element) {
+            return <Element {...props}>{children}</Element>;
         }
+        return <pre ref={this.setEl}><code className={className}>{children}</code></pre>;
     }
 }
 
diff --git a/src/optimized.js b/src/optimized.js
index d5fe955..e97f670 100644
--- a/src/optimized.js
+++ b/src/optimized.js
@@ -1,6 +1,5 @@
 import hljs from'highlight.js/lib/highlight';
 import React from'react';
-import ReactDOM from'react-dom';
 
 class Highlight extends React.Component {
     componentDidMount() {
@@ -13,8 +12,7 @@ class Highlight extends React.Component {
 
     highlightCode() {
         const {className, languages} = this.props;
-        const domNode = ReactDOM.findDOMNode(this);
-        const nodes = domNode.querySelectorAll('pre code');
+        const nodes = this.el.querySelectorAll('pre code');
 
         if ((languages.length === 0) && className) {
             languages.push(className);
@@ -24,29 +22,31 @@ class Highlight extends React.Component {
             hljs.registerLanguage(lang, require('highlight.js/lib/languages/' + lang));
         });
 
-        let i;
-        for (i = 0; i < nodes.length; i++) {
-            hljs.highlightBlock(nodes[i]);
+        for (let i = 0; i < nodes.length; i++) {
+            hljs.highlightBlock(nodes[i])
         }
     }
 
+    setEl = (el) => {
+        this.el = el;
+    };
+
     render() {
-        const {children, className, element, innerHTML} = this.props;
-        let Element = element ? React.DOM[element] : null;
+        const {children, className, element: Element, innerHTML} = this.props;
+        const props = { ref: this.setEl, className };
 
         if (innerHTML) {
-            if (!Element) {
-                Element = React.DOM.div
-            }
-
-            return Element({dangerouslySetInnerHTML: {__html: children}, className: className || null}, null);
-        } else {
+            props.dangerouslySetInnerHTML = { __html: children };
             if (Element) {
-                return Element({className}, children);
-            } else {
-                return <pre><code className={className}>{children}</code></pre>;
+                return <Element {...props} />;
             }
+            return <div {...props} />;
+        }
+
+        if (Element) {
+            return <Element {...props}>{children}</Element>;
         }
+        return <pre ref={this.setEl}><code className={className}>{children}</code></pre>;
     }
 }