diff --git a/composer.json b/composer.json
index f3f2bec..3480145 100644
--- a/composer.json
+++ b/composer.json
@@ -14,7 +14,8 @@
 	"require": {
 		"php": ">=5.5.0",
 		"benestar/asparagus": "~0.4",
-		"monolog/monolog": "~1.18"
+		"monolog/monolog": "~1.18",
+		"wikibase/data-model": "~6.0"
 	},
 	"autoload": {
 		"psr-4": {
diff --git a/src/EntityImporter.php b/src/EntityImporter.php
index 299cef5..2545f64 100644
--- a/src/EntityImporter.php
+++ b/src/EntityImporter.php
@@ -5,13 +5,11 @@
 use Psr\Log\LoggerInterface;
 use User;
 use Wikibase\DataModel\Entity\BasicEntityIdParser;
-use Wikibase\DataModel\Entity\Entity;
+use Wikibase\DataModel\Entity\EntityDocument;
 use Wikibase\DataModel\Entity\EntityId;
 use Wikibase\DataModel\Entity\EntityIdParsingException;
 use Wikibase\DataModel\Entity\EntityIdValue;
 use Wikibase\DataModel\Entity\Item;
-use Wikibase\DataModel\Entity\ItemId;
-use Wikibase\DataModel\Serializers\StatementSerializer;
 use Wikibase\DataModel\Snak\PropertyValueSnak;
 use Wikibase\DataModel\Statement\StatementList;
 use Wikibase\Repo\Store\WikiPageEntityStore;
@@ -135,7 +133,7 @@ private function importBatch( array $batch ) {
 		return $stashedEntities;
 	}
 
-	private function createEntity( Entity $entity ) {
+	private function createEntity( EntityDocument $entity ) {
 		$entity->setId( null );
 
 		$entity->setStatements( new StatementList() );
@@ -171,7 +169,7 @@ private function getBadgeItems( array $entities ) {
 		return $badgeItems;
 	}
 
-	private function getReferencedEntities( Entity $entity ) {
+	private function getReferencedEntities( EntityDocument $entity ) {
 		$snaks = $entity->getStatements()->getAllSnaks();
 		$entities = array();
 
diff --git a/src/StatementsImporter.php b/src/StatementsImporter.php
index 27d2aa9..656fbeb 100644
--- a/src/StatementsImporter.php
+++ b/src/StatementsImporter.php
@@ -3,24 +3,15 @@
 namespace Wikibase\Import;
 
 use ApiMain;
-use DataValues\Serializers\DataValueSerializer;
 use Serializers\Serializer;
 use FauxRequest;
 use Psr\Log\LoggerInterface;
 use RequestContext;
 use User;
-use Wikibase\DataModel\DeserializerFactory;
 use Wikibase\DataModel\Entity\BasicEntityIdParser;
-use Wikibase\DataModel\Entity\Entity;
+use Wikibase\DataModel\Entity\EntityDocument;
 use Wikibase\DataModel\Entity\EntityId;
-use Wikibase\DataModel\Entity\EntityIdValue;
-use Wikibase\DataModel\Entity\Item;
-use Wikibase\DataModel\Entity\ItemId;
-use Wikibase\DataModel\Entity\PropertyId;
 use Wikibase\DataModel\Serializers\StatementSerializer;
-use Wikibase\DataModel\Snak\PropertyNoValueSnak;
-use Wikibase\DataModel\Snak\PropertySomeValueSnak;
-use Wikibase\DataModel\Snak\PropertyValueSnak;
 use Wikibase\DataModel\Statement\Statement;
 use Wikibase\DataModel\Statement\StatementList;
 
@@ -52,7 +43,7 @@ public function __construct(
 		$this->idParser = new BasicEntityIdParser();
 	}
 
-	public function importStatements( Entity $entity ) {
+	public function importStatements( EntityDocument $entity ) {
 		$statements = $entity->getStatements();
 
 		$this->logger->info( 'Adding statements: ' . $entity->getId()->getSerialization() );